From 5aaaa5a709c436186852c09f4cbf7645bf20014f Mon Sep 17 00:00:00 2001
From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
Date: Mon, 6 Jul 2026 16:31:46 -0700
Subject: [PATCH] [nvbugs/6414760][fix] Port legacy Qwen cpprunner e2e test to
 PyTorch LLM API

Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
---
 tests/integration/defs/test_e2e.py | 73 ++++++++----------------------
 1 file changed, 20 insertions(+), 53 deletions(-)

diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py
index 11c08a284c74..3769f51ec40d 100644
--- a/tests/integration/defs/test_e2e.py
+++ b/tests/integration/defs/test_e2e.py
@@ -27,7 +27,7 @@
 from defs.trt_test_alternative import (check_call, check_call_negative_test,
                                        check_output, print_info, print_warning)
 
-from .common import convert_weights, get_mmlu_accuracy, venv_check_call
+from .common import get_mmlu_accuracy, venv_check_call
 from .conftest import (get_device_count, get_sm_version, llm_models_root,
                        skip_no_sm120, skip_nvlink_inactive, skip_post_blackwell,
                        skip_pre_ada, skip_pre_blackwell, skip_pre_hopper,
@@ -72,62 +72,29 @@ def test_gpt3_175b_1layers_build_only(llm_root, llm_venv, engine_dir):
 @pytest.mark.parametrize("model_name,model_path", [
     ("DeepSeek-R1-Distill-Qwen-1.5B", "DeepSeek-R1-Distill-Qwen-1.5B"),
 ])
-def test_qwen_e2e_cpprunner_large_new_tokens(model_name, model_path, llm_venv,
-                                             qwen_example_root, cmodel_dir,
-                                             engine_dir):
-    """RCCA: https://nvbugs/5238105"""
-    model_dir = convert_weights(
-        llm_venv=llm_venv,
-        example_root=qwen_example_root,
-        cmodel_dir=cmodel_dir,
-        model=model_name,
-        model_path=f"{llm_models_root()}/{model_path}",
-    )
-
-    build_cmd = [
-        "trtllm-build", f"--checkpoint_dir={model_dir}",
-        f"--output_dir={engine_dir}", "--gemm_plugin=float16",
-        "--max_num_tokens=32768"
-    ]
+def test_qwen_e2e_cpprunner_large_new_tokens(model_name, model_path, llm_venv):
+    """RCCA: https://nvbugs/5238105 - none of n>1 sequences may be empty."""
+    from tensorrt_llm import LLM, SamplingParams
 
-    check_call(" ".join(build_cmd), shell=True, env=llm_venv._new_env)
+    prompt = r"<｜begin▁of▁sentence｜><｜User｜>The operation $\otimes$ is defined for all nonzero numbers by $a \otimes b = \frac{a^{2}}{b}$. Determine $[(1 \otimes 2) \otimes 3] - [1 \otimes (2 \otimes 3)]$. Let's think step by step and output the final answer within \boxed{}.<｜Assistant｜>"
 
-    from transformers import AutoTokenizer
+    sampling_params = SamplingParams(
+        max_tokens=1024,
+        n=4,
+        temperature=0.6,
+        top_p=1.0,
+        top_k=1024,
+    )
 
-    from tensorrt_llm.runtime import PYTHON_BINDINGS
+    with LLM(model=f"{llm_models_root()}/{model_path}",
+             max_batch_size=8,
+             max_seq_len=4224) as llm:
+        outputs = llm.generate([prompt], sampling_params=sampling_params)
 
-    if PYTHON_BINDINGS:
-        from tensorrt_llm.runtime import ModelRunnerCpp
-    tokenizer = AutoTokenizer.from_pretrained(
-        f"{llm_models_root()}/{model_path}",
-        trust_remote_code=True,
-        use_fast=False)
-
-    message = r"<｜begin▁of▁sentence｜><｜User｜>The operation $\otimes$ is defined for all nonzero numbers by $a \otimes b = \frac{a^{2}}{b}$. Determine $[(1 \otimes 2) \otimes 3] - [1 \otimes (2 \otimes 3)]$. Let's think step by step and output the final answer within \boxed{}.<｜Assistant｜>"
-
-    inputs = tokenizer(message, return_tensors='pt',
-                       add_special_tokens=False)['input_ids']
-
-    runner = ModelRunnerCpp.from_dir(engine_dir=f"{engine_dir}",
-                                     max_input_len=128,
-                                     max_output_len=4096,
-                                     max_batch_size=8)
-
-    outputs = runner.generate(inputs,
-                              end_id=tokenizer.eos_token_id,
-                              pad_id=tokenizer.pad_token_id,
-                              temperature=0.6,
-                              top_p=1.0,
-                              top_k=1024,
-                              max_new_tokens=1024,
-                              return_dict=True,
-                              min_length=1,
-                              num_return_sequences=4,
-                              output_sequence_lengths=True)
-
-    seq_lengths = outputs['sequence_lengths']
-    assert not (seq_lengths == 0).any(
-    ), f"Found zero length in sequence_lengths tensor: {seq_lengths}"
+    completions = outputs[0].outputs
+    seq_lengths = [len(c.token_ids) for c in completions]
+    assert all(length > 0 for length in seq_lengths), \
+        f"Found zero-length completion: {seq_lengths}"
 
 
 # TODO replace the trtllm_bench_prolog