quickstart_advanced.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import json
import time
from typing import Literal

from tensorrt_llm import LLM, SamplingParams
from tensorrt_llm.llmapi import (AttentionDpConfig, AutoDecodingConfig,
                                 CudaGraphConfig, DFlashDecodingConfig,
                                 DraftTargetDecodingConfig,
                                 Eagle3DecodingConfig, KvCacheConfig, MoeConfig,
                                 MTPDecodingConfig, NGramDecodingConfig,
                                 TorchCompileConfig)

example_prompts = [
    "Hello, my name is",
    "The capital of France is",
    "The future of AI is",
]


def _parse_kv_cache_manager_v2(value: str) -> bool | Literal["auto"]:
    if value == "auto":
        return "auto"
    if value == "true":
        return True
    if value == "false":
        return False
    raise argparse.ArgumentTypeError("expected one of: auto, true, false")


def add_llm_args(parser):
    parser.add_argument('--model_dir',
                        type=str,
                        required=True,
                        help="Model checkpoint directory.")
    parser.add_argument("--prompt",
                        type=str,
                        nargs="+",
                        help="A single or a list of text prompts.")
    parser.add_argument('--checkpoint_format',
                        type=str,
                        default=None,
                        choices=["HF", "mistral"],
                        help="Model checkpoint format.")
    # Build config
    parser.add_argument("--max_seq_len",
                        type=int,
                        default=None,
                        help="The maximum sequence length.")
    parser.add_argument("--max_batch_size",
                        type=int,
                        default=2048,
                        help="The maximum batch size.")
    parser.add_argument(
        "--max_num_tokens",
        type=int,
        default=8192,
        help=
        "The maximum total tokens (context + generation) across all sequences in a batch."
    )

    # Parallelism
    parser.add_argument('--attention_backend',
                        type=str,
                        default='TRTLLM',
                        choices=['VANILLA', 'TRTLLM', 'FLASHINFER'])
    parser.add_argument(
        '--moe_backend',
        type=str,
        default='AUTO',
        choices=[
            'AUTO', 'CUTLASS', 'TRTLLM', 'VANILLA', 'DEEPGEMM', 'CUTEDSL',
            'TRITON'
        ],
        help=
        'MoE backend to use. AUTO selects default backend based on model. It currently doesn\'t always give the best choice for all scenarios. The capabilities of auto selection will be improved in future releases.'
    )
    parser.add_argument('--enable_attention_dp',
                        default=False,
                        action='store_true')
    parser.add_argument('--attention_dp_enable_balance',
                        default=False,
                        action='store_true')
    parser.add_argument('--attention_dp_time_out_iters', type=int, default=0)
    parser.add_argument('--attention_dp_batching_wait_iters',
                        type=int,
                        default=0)
    parser.add_argument('--tp_size', type=int, default=1)
    parser.add_argument('--pp_size', type=int, default=1)
    parser.add_argument('--orchestrator_type',
                        type=str,
                        default=None,
                        choices=[None, 'rpc', 'ray'],
                        help='Orchestrator type for multi-GPU execution')
    parser.add_argument('--moe_ep_size', type=int, default=-1)
    parser.add_argument('--moe_tp_size', type=int, default=-1)
    parser.add_argument('--moe_cluster_size', type=int, default=-1)
    parser.add_argument(
        '--use_low_precision_moe_combine',
        default=False,
        action='store_true',
        help='Use low precision combine in MoE (only for NVFP4 quantization)')
    parser.add_argument(
        '--moe_load_balancer_config',
        type=str,
        default=None,
        help='Path to a YAML file for MoE load balancer (EPLB) configuration.')

    # KV cache
    parser.add_argument('--kv_cache_dtype', type=str, default='auto')
    parser.add_argument('--disable_kv_cache_reuse',
                        default=False,
                        action='store_true')
    parser.add_argument("--tokens_per_block", type=int, default=32)
    parser.add_argument('--mamba_ssm_cache_dtype',
                        type=str,
                        default='bfloat16',
                        choices=['auto', 'float16', 'bfloat16', 'float32'],
                        help='Data type for Mamba SSM cache.')
    parser.add_argument(
        '--mamba_ssm_stochastic_rounding',
        default=False,
        action='store_true',
        help=
        'Enable stochastic rounding for Mamba SSM state updates (fp16 only, FlashInfer limitation).'
    )
    parser.add_argument(
        '--mamba_ssm_philox_rounds',
        type=int,
        default=10,
        help=
        'Number of Philox rounds for stochastic rounding PRNG (default: 10). Higher values give better randomness.'
    )
    parser.add_argument('--log_kv_cache_events',
                        default=False,
                        action='store_true')
    parser.add_argument(
        '--use_kv_cache_manager_v2',
        default='auto',
        type=_parse_kv_cache_manager_v2,
        metavar='{auto,true,false}',
        help=
        'Whether to use KVCacheManagerV2 for KV cache management (PyTorch backend). Defaults to model-specific auto selection.',
    )

    # Runtime
    parser.add_argument('--disable_overlap_scheduler',
                        default=False,
                        action='store_true')
    parser.add_argument('--enable_chunked_prefill',
                        default=False,
                        action='store_true')
    parser.add_argument('--use_cuda_graph', default=False, action='store_true')
    parser.add_argument('--cuda_graph_padding_enabled',
                        default=False,
                        action='store_true')
    parser.add_argument('--cuda_graph_batch_sizes',
                        nargs='+',
                        type=int,
                        default=None)
    parser.add_argument('--print_iter_log',
                        default=False,
                        action='store_true',
                        help='Print iteration logs during execution')
    parser.add_argument('--use_torch_compile',
                        default=False,
                        action='store_true',
                        help='Use torch.compile to optimize the model')
    parser.add_argument('--use_piecewise_cuda_graph',
                        default=False,
                        action='store_true',
                        help='Use piecewise CUDA graph to optimize the model')
    parser.add_argument('--apply_chat_template',
                        default=False,
                        action='store_true')
    parser.add_argument('--custom_tokenizer',
                        type=str,
                        default=None,
                        help='Override the tokenizer. Accepts a built-in alias '
                        " or a fully-qualified class import path.")

    # Sampling
    parser.add_argument("--max_tokens", type=int, default=64)
    parser.add_argument("--temperature", type=float, default=None)
    parser.add_argument("--top_k", type=int, default=None)
    parser.add_argument("--top_p", type=float, default=None)
    parser.add_argument('--load_format', type=str, default='auto')
    parser.add_argument('--n', type=int, default=1)
    parser.add_argument('--best_of', type=int, default=None)
    parser.add_argument('--max_beam_width', type=int, default=1)

    # Speculative decoding
    parser.add_argument('--spec_decode_algo', type=str, default=None)
    parser.add_argument('--spec_decode_max_draft_len', type=int, default=1)
    parser.add_argument('--draft_model_dir', type=str, default=None)
    parser.add_argument('--max_matching_ngram_size', type=int, default=5)
    parser.add_argument('--use_dynamic_tree',
                        default=False,
                        action='store_true')
    parser.add_argument('--dynamic_tree_max_topK', type=int, default=None)
    parser.add_argument('--eagle3_model_arch',
                        type=str,
                        default="llama3",
                        choices=["llama3", "mistral_large3"],
                        help="The model architecture of the eagle3 model.")
    parser.add_argument('--max_total_draft_tokens', type=int, default=None)

    # Relaxed acceptance
    parser.add_argument('--use_relaxed_acceptance_for_thinking',
                        default=False,
                        action='store_true')
    parser.add_argument('--relaxed_topk', type=int, default=1)
    parser.add_argument('--relaxed_delta', type=float, default=0.)

    # HF
    parser.add_argument('--trust_remote_code',
                        default=False,
                        action='store_true')
    parser.add_argument('--return_context_logits',
                        default=False,
                        action='store_true')
    parser.add_argument('--return_generation_logits',
                        default=False,
                        action='store_true')
    parser.add_argument('--prompt_logprobs', type=int, default=None)
    parser.add_argument('--logprobs', type=int, default=None)

    parser.add_argument('--additional_model_outputs',
                        type=str,
                        default=None,
                        nargs='+')

    return parser


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="LLM models with the PyTorch workflow.")
    parser = add_llm_args(parser)
    parser.add_argument("--kv_cache_fraction", type=float, default=0.9)
    parser.add_argument(
        "--streaming",
        action="store_true",
        default=False,
        help="Use streaming generate_async instead of generate.")
    parser.add_argument(
        "--concurrency",
        type=int,
        default=1,
        help="Number of concurrent requests to submit simultaneously.")
    parser.add_argument(
        "--dataset",
        type=str,
        default=None,
        help=
        "Path to a JSONL dataset file. Each line: {\"question_id\": N, \"question\": [\"prompt\"]}."
    )
    parser.add_argument(
        "--num_samples",
        type=int,
        default=None,
        required=False,
        help="Limit number of samples from dataset (for fast debugging).")
    args = parser.parse_args()
    return args


def setup_llm(args, **kwargs):
    kv_cache_config = KvCacheConfig(
        enable_block_reuse=not args.disable_kv_cache_reuse,
        free_gpu_memory_fraction=args.kv_cache_fraction,
        dtype=args.kv_cache_dtype,
        tokens_per_block=args.tokens_per_block,
        use_kv_cache_manager_v2=args.use_kv_cache_manager_v2,
        mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
        mamba_ssm_stochastic_rounding=args.mamba_ssm_stochastic_rounding,
        mamba_ssm_philox_rounds=args.mamba_ssm_philox_rounds,
        event_buffer_max_size=1024 if args.log_kv_cache_events else 0)

    spec_decode_algo = args.spec_decode_algo.upper(
    ) if args.spec_decode_algo is not None else None

    if spec_decode_algo == 'MTP':
        speculative_model = (args.draft_model_dir if args.draft_model_dir
                             is not None else args.model_dir)
        spec_config = MTPDecodingConfig(
            max_draft_len=args.spec_decode_max_draft_len,
            use_relaxed_acceptance_for_thinking=args.
            use_relaxed_acceptance_for_thinking,
            relaxed_topk=args.relaxed_topk,
            relaxed_delta=args.relaxed_delta,
            use_dynamic_tree=args.use_dynamic_tree,
            dynamic_tree_max_topK=args.dynamic_tree_max_topK,
            max_total_draft_tokens=args.max_total_draft_tokens,
            speculative_model=speculative_model)
    elif spec_decode_algo == "EAGLE3":
        spec_config = Eagle3DecodingConfig(
            max_draft_len=args.spec_decode_max_draft_len,
            speculative_model=args.draft_model_dir,
            use_dynamic_tree=args.use_dynamic_tree,
            dynamic_tree_max_topK=args.dynamic_tree_max_topK,
            eagle3_model_arch=args.eagle3_model_arch,
            max_total_draft_tokens=args.max_total_draft_tokens)
    elif spec_decode_algo == "DFLASH":
        spec_config = DFlashDecodingConfig(
            max_draft_len=args.spec_decode_max_draft_len,
            speculative_model=args.draft_model_dir)
    elif spec_decode_algo == "DRAFT_TARGET":
        spec_config = DraftTargetDecodingConfig(
            max_draft_len=args.spec_decode_max_draft_len,
            speculative_model=args.draft_model_dir)
    elif spec_decode_algo == "NGRAM":
        spec_config = NGramDecodingConfig(
            max_draft_len=args.spec_decode_max_draft_len,
            max_matching_ngram_size=args.max_matching_ngram_size,
            is_keep_all=True,
            is_use_oldest=True,
            is_public_pool=True,
        )
    elif spec_decode_algo == "AUTO":
        spec_config = AutoDecodingConfig()
    else:
        spec_config = None

    cuda_graph_config = CudaGraphConfig(
        batch_sizes=args.cuda_graph_batch_sizes,
        enable_padding=args.cuda_graph_padding_enabled,
    ) if args.use_cuda_graph else None

    attention_dp_config = AttentionDpConfig(
        enable_balance=args.attention_dp_enable_balance,
        timeout_iters=args.attention_dp_time_out_iters,
        batching_wait_iters=args.attention_dp_batching_wait_iters,
    )

    llm = LLM(
        model=args.model_dir,
        backend='pytorch',
        checkpoint_format=args.checkpoint_format,
        disable_overlap_scheduler=args.disable_overlap_scheduler,
        kv_cache_config=kv_cache_config,
        attn_backend=args.attention_backend,
        cuda_graph_config=cuda_graph_config,
        load_format=args.load_format,
        print_iter_log=args.print_iter_log,
        enable_iter_perf_stats=args.print_iter_log,
        torch_compile_config=TorchCompileConfig(
            enable_fullgraph=args.use_torch_compile,
            enable_piecewise_cuda_graph= \
                args.use_piecewise_cuda_graph)
        if args.use_torch_compile else None,
        moe_config=MoeConfig(backend=args.moe_backend, use_low_precision_moe_combine=args.use_low_precision_moe_combine, load_balancer=args.moe_load_balancer_config),
        max_seq_len=args.max_seq_len,
        max_batch_size=args.max_batch_size,
        max_num_tokens=args.max_num_tokens,
        enable_attention_dp=args.enable_attention_dp,
        attention_dp_config=attention_dp_config,
        tensor_parallel_size=args.tp_size,
        pipeline_parallel_size=args.pp_size,
        moe_expert_parallel_size=args.moe_ep_size,
        moe_tensor_parallel_size=args.moe_tp_size,
        moe_cluster_parallel_size=args.moe_cluster_size,
        enable_chunked_prefill=args.enable_chunked_prefill,
        speculative_config=spec_config,
        trust_remote_code=args.trust_remote_code,
        gather_generation_logits=args.return_generation_logits,
        max_beam_width=args.max_beam_width,
        orchestrator_type=args.orchestrator_type,
        custom_tokenizer=args.custom_tokenizer,
        **kwargs)

    use_beam_search = args.max_beam_width > 1
    best_of = args.best_of or args.n
    if use_beam_search:
        if args.n == 1 and args.best_of is None:
            args.n = args.max_beam_width
        assert best_of <= args.max_beam_width, f"beam width: {best_of}, should be less or equal to max_beam_width: {args.max_beam_width}"

    assert best_of >= args.n, f"In sampling mode best_of value: {best_of} should be greater than or equal to n: {args.n}"

    sampling_params = SamplingParams(
        max_tokens=args.max_tokens,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        return_context_logits=args.return_context_logits,
        return_generation_logits=args.return_generation_logits,
        logprobs=args.logprobs,
        prompt_logprobs=args.prompt_logprobs,
        n=args.n,
        best_of=best_of,
        use_beam_search=use_beam_search,
        additional_model_outputs=args.additional_model_outputs,
    )
    return llm, sampling_params


def main():
    args = parse_arguments()

    if args.dataset:
        prompts = []
        with open(args.dataset, 'r') as f:
            for line in f:
                entry = json.loads(line.strip())
                prompts.append(entry["question"][-1])
        if args.num_samples is not None:
            prompts = prompts[:args.num_samples]
    else:
        prompts = args.prompt if args.prompt else example_prompts

    llm, sampling_params = setup_llm(args)
    new_prompts = []
    if args.apply_chat_template:
        for prompt in prompts:
            messages = [{"role": "user", "content": f"{prompt}"}]
            new_prompts.append(
                llm.tokenizer.apply_chat_template(messages,
                                                  tokenize=False,
                                                  add_generation_prompt=True))
        prompts = new_prompts
    llm.generate(prompts, sampling_params)

    accept_rates = []
    total_tokens = 0
    total_iterations = 0

    if args.streaming:
        for i, prompt in enumerate(prompts):
            num_tokens = 0
            num_iterations = 0
            for output in llm.generate_async(prompt,
                                             sampling_params,
                                             streaming=True):
                new_tokens = output.outputs[0].token_ids
                num_tokens = len(new_tokens)
                num_iterations += 1
            if num_iterations > 0:
                accept_rate = num_tokens / num_iterations
                accept_rates.append(accept_rate)
                total_tokens += num_tokens
                total_iterations += num_iterations
                print(f"[{i}] Accept rate: {accept_rate:.2f} "
                      f"(tokens={num_tokens}, iterations={num_iterations})")
            generated_text = output.outputs[0].text
            print(
                f"[{i}] Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..."
            )

        if accept_rates:
            avg_accept_rate = sum(accept_rates) / len(accept_rates)
            global_accept_rate = total_tokens / total_iterations
            print("\n=== Accept Rate Summary ===")
            print(f"Total prompts: {len(accept_rates)}")
            print(f"Mean accept rate (per-prompt avg): {avg_accept_rate:.2f}")
            print(
                f"Global accept rate (total_tokens/total_iters): {global_accept_rate:.2f}"
            )
            print(
                f"Total tokens: {total_tokens}, Total iterations: {total_iterations}"
            )
        return

    concurrency = args.concurrency

    if concurrency > 1:
        # Submit all prompts concurrently in batches of `concurrency`
        for batch_start in range(0, len(prompts), concurrency):
            batch_end = min(batch_start + concurrency, len(prompts))
            batch_prompts = prompts[batch_start:batch_end]

            futures = []
            for prompt in batch_prompts:
                future = llm.generate_async(prompt,
                                            sampling_params,
                                            streaming=True)
                futures.append(future)

            for j, future in enumerate(futures):
                i = batch_start + j
                num_tokens = 0
                num_iterations = 0
                for output in future:
                    new_tokens = output.outputs[0].token_ids
                    num_tokens = len(new_tokens)
                    num_iterations += 1
                if num_iterations > 0:
                    accept_rate = num_tokens / num_iterations
                    accept_rates.append(accept_rate)
                    total_tokens += num_tokens
                    total_iterations += num_iterations
                    print(f"[{i}] Accept rate: {accept_rate:.2f} "
                          f"(tokens={num_tokens}, iterations={num_iterations})")
                generated_text = output.outputs[0].text
                print(
                    f"[{i}] Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..."
                )
    else:
        # Use streaming generate_async to count decode iterations for accept rate
        for i, prompt in enumerate(prompts):
            num_tokens = 0
            num_iterations = 0
            for output in llm.generate_async(prompt,
                                             sampling_params,
                                             streaming=True):
                new_tokens = output.outputs[0].token_ids
                num_tokens = len(new_tokens)
                num_iterations += 1
            if num_iterations > 0:
                accept_rate = num_tokens / num_iterations
                accept_rates.append(accept_rate)
                total_tokens += num_tokens
                total_iterations += num_iterations
                print(f"[{i}] Accept rate: {accept_rate:.2f} "
                      f"(tokens={num_tokens}, iterations={num_iterations})")
            for sequence_idx, sequence in enumerate(output.outputs):
                generated_text = sequence.text
                sequence_id_text = f"[{sequence_idx}]" if args.max_beam_width > 1 or args.n > 1 else ""
                print(
                    f"[{i}]{sequence_id_text} Prompt: {prompt[:80]!r}..., Generated text: {generated_text[:200]!r}..."
                )
                if args.return_context_logits:
                    print(
                        f"[{i}]{sequence_id_text} Context logits: {output.context_logits}"
                    )
                if args.return_generation_logits:
                    print(
                        f"[{i}]{sequence_id_text} Generation logits: {sequence.generation_logits}"
                    )
                if args.prompt_logprobs is not None:
                    print(
                        f"[{i}]{sequence_id_text} Prompt logprobs: {sequence.prompt_logprobs}"
                    )
                if args.logprobs is not None:
                    print(
                        f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}"
                    )

                if args.additional_model_outputs:
                    for output_name in args.additional_model_outputs:
                        if sequence.additional_context_outputs:
                            print(
                                f"[{i}]{sequence_id_text} Context {output_name}: {sequence.additional_context_outputs[output_name]}"
                            )
                        print(
                            f"[{i}]{sequence_id_text} Generation {output_name}: {sequence.additional_generation_outputs[output_name]}"
                        )

    if args.log_kv_cache_events:
        time.sleep(1)  # Wait for events to be dispatched
        events = llm.get_kv_cache_events(5)
        print("=== KV_CACHE_EVENTS_START ===")
        print(json.dumps(events, indent=2))
        print("=== KV_CACHE_EVENTS_END ===")


if __name__ == '__main__':
    main()