From 54f70310515ab74c4521fc2edcea519c2476fdb4 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:46:11 +0800 Subject: [PATCH 01/16] [None][chore] Remove all Llama-3.1-8B test coverage across integration and unit tests Fully removes Llama-3.1-8B / Llama-3.1-8B-Instruct from the test suite, including dedicated accuracy test classes, shared multi-model parametrizations in test_e2e.py/disaggregated/examples/llmapi tests, corresponding test-list entries (waives.txt, qa/*, test-db/*), accuracy reference YAMLs, the RocketKV sparse-attention unit test, spec-decoding and AutoDeploy smoke tests, and the Llama architecture unit test (test_modeling_llama.py, with its shared Scenario/reduce_llama_config helpers preserved for test_modeling_mllama.py). Files left with no remaining tests were deleted outright. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/defs/conftest.py | 12 - .../defs/disaggregated/test_ad_disagg.py | 73 -- ...isagg_config_ctxtp2_gentp2_llama31_8b.yaml | 18 - .../defs/disaggregated/test_disaggregated.py | 205 ---- .../test_disaggregated_single_gpu.py | 115 -- .../examples/test_ad_speculative_decoding.py | 1075 +---------------- .../defs/llmapi/test_llm_examples.py | 3 - tests/integration/defs/test_e2e.py | 404 +------ .../test_lists/qa/llm_function_core.txt | 6 +- .../test_lists/qa/llm_function_multinode.txt | 3 - .../test_lists/qa/llm_spark_func.yml | 4 - .../integration/test_lists/test-db/l0_a10.yml | 1 - .../test_lists/test-db/l0_b200.yml | 6 - .../test_lists/test-db/l0_dgx_h200.yml | 1 - .../test_lists/test-db/l0_h100.yml | 12 - .../test_lists/test-db/l0_rtx_pro_6000.yml | 3 - tests/integration/test_lists/waives.txt | 16 +- .../sparse/rocketkv/test_rocketkv.py | 88 +- .../_torch/modeling/test_modeling_llama.py | 714 ----------- .../hw_agnostic/test_kv_cache_reuse.py | 84 -- .../speculative/hw_agnostic/test_ngram.py | 68 -- .../speculative/hw_agnostic/test_pard.py | 115 -- .../_torch/speculative/hw_agnostic/test_sa.py | 205 +--- .../speculative/hw_agnostic/test_spec_gate.py | 135 --- .../hw_agnostic/test_user_provided.py | 90 -- .../_torch/speculative/test_eagle3.py | 608 +--------- .../_utils_test/_model_test_utils.py | 11 +- .../smoke/test_ad_allreduce_strategies.py | 307 ----- .../smoke/test_ad_build_small_multi.py | 49 - .../singlegpu/models/test_eagle.py | 144 +-- .../smoke/test_ad_build_small_single.py | 26 - .../smoke/test_ad_guided_decoding_regex.py | 73 -- .../smoke/test_ad_speculative_decoding.py | 160 +-- .../singlegpu/smoke/test_ad_trtllm_serve.py | 117 -- .../singlegpu/smoke/test_disagg.py | 59 - .../apps/_test_disagg_serving_multi_nodes.py | 213 ---- ...g_serving_multi_nodes_service_discovery.py | 223 ---- .../apps/_test_openai_chat_guided_decoding.py | 5 +- .../llmapi/apps/_test_openai_multi_nodes.py | 257 ---- .../apps/_test_trtllm_serve_benchmark.py | 8 +- tests/unittest/llmapi/test_llm.py | 1 - tests/unittest/llmapi/test_llm_pytorch.py | 33 +- tests/unittest/llmapi/test_session_reuse.py | 8 +- tests/unittest/metrics/test_collector.py | 4 +- .../scripts/test_check_model_registry.py | 4 +- 45 files changed, 38 insertions(+), 5728 deletions(-) delete mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_llama31_8b.yaml delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_pard.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py delete mode 100644 tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py delete mode 100644 tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py delete mode 100644 tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py delete mode 100644 tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py delete mode 100644 tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py delete mode 100644 tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py delete mode 100644 tests/unittest/llmapi/apps/_test_openai_multi_nodes.py diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index bcea8e086953..2f0e913a83e0 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -680,18 +680,6 @@ def llama_model_root(request): if request.param == "TinyLlama-1.1B-Chat-v1.0": llama_model_root = os.path.join(models_root, "llama-models-v2", "TinyLlama-1.1B-Chat-v1.0") - elif request.param == "llama-3.1-8b": - llama_model_root = os.path.join(models_root, "llama-3.1-model", - "Meta-Llama-3.1-8B") - elif request.param == "llama-3.1-8b-instruct-hf-fp8": - llama_model_root = os.path.join(models_root, "llama-3.1-model", - "Llama-3.1-8B-Instruct-FP8") - elif request.param == "llama-3.1-8b-instruct": - llama_model_root = os.path.join(models_root, "llama-3.1-model", - "Llama-3.1-8B-Instruct") - elif request.param == "llama-3.1-8b-hf-nvfp4": - llama_model_root = os.path.join(models_root, "nvfp4-quantized", - "Meta-Llama-3.1-8B") assert os.path.exists( llama_model_root ), f"{llama_model_root} does not exist under NFS LLM_MODELS_ROOT dir" diff --git a/tests/integration/defs/disaggregated/test_ad_disagg.py b/tests/integration/defs/disaggregated/test_ad_disagg.py index 9d8bff8a3d42..b4d49e845673 100644 --- a/tests/integration/defs/disaggregated/test_ad_disagg.py +++ b/tests/integration/defs/disaggregated/test_ad_disagg.py @@ -33,7 +33,6 @@ from tensorrt_llm import DisaggregatedParams, SamplingParams from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM from tensorrt_llm._utils import set_mpi_comm -from tensorrt_llm.llmapi import Eagle3DecodingConfig cloudpickle.register_pickle_by_value(sys.modules[__name__]) MPI.pickle.__init__( @@ -67,30 +66,9 @@ def skip_b300(): AUTODEPLOY_DISAGG_SEED = 1234 REDUCED_TINYLLAMA_LAYERS = 2 REDUCED_DEEPSEEK_LAYERS = 2 -LLAMA_EAGLE3_EXPECTED_TEXT = " Berlin\nWhat is the capital of France? Paris\nWhat is the capital of" -LLAMA_EAGLE3_EXPECTED_TOKEN_IDS = [ - 20437, - 198, - 3923, - 374, - 279, - 6864, - 315, - 9822, - 30, - 12366, - 198, - 3923, - 374, - 279, - 6864, - 315, -] MODEL_PATHS = { - "EAGLE3-LLaMA3.1-Instruct-8B": "EAGLE3-LLaMA3.1-Instruct-8B", - "Llama-3.1-8B-Instruct": "llama-3.1-model/Llama-3.1-8B-Instruct/", "TinyLlama-1.1B-Chat-v1.0": "llama-models-v2/TinyLlama-1.1B-Chat-v1.0", "DeepSeek-V3-Lite": "DeepSeek-V3-Lite/bf16", } @@ -657,20 +635,6 @@ def test_chunked_prefill_handoff(model): # --------------------------------------------------------------------------- -def llama_eagle3_config(): - return { - "speculative_config": Eagle3DecodingConfig( - max_draft_len=3, - speculative_model=model_path("EAGLE3-LLaMA3.1-Instruct-8B"), - eagle3_one_model=True, - eagle3_layers_to_capture={1, 15, 28}, - ), - # Force the Eagle3 draft to match the BF16 Llama 3.1 target. Shared KV - # cache management requires matching target and draft KV dtypes. - "speculative_model_kwargs": {"torch_dtype": "bfloat16"}, - } - - def get_ucx_tls() -> str: """Get UCX_TLS value based on GPU architecture. @@ -1032,40 +996,3 @@ def test_async_sharded_generation_handoff(): assert outputs["context"].token_ids == aggregate_output.token_ids[:1] assert outputs["generation"].text == aggregate_output.text assert outputs["generation"].token_ids == aggregate_output.token_ids - - -@skip_pre_hopper -@pytest.mark.threadleak(enabled=False) -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.skip_less_device(2) -@pytest.mark.timeout(900) -def test_async_eagle3_full_model_handoff(): - sampling_params_kwargs = { - "max_tokens": 16, - "ignore_eos": True, - "top_k": 1, - "seed": AUTODEPLOY_DISAGG_SEED, - } - extra_config = llama_eagle3_config() - outputs = run_context_then_generation_handoff( - "Llama-3.1-8B-Instruct", - worker_world_sizes=(1, 1), - generation_overlap=True, - prompt="What is the capital of Germany?", - sampling_params_kwargs=sampling_params_kwargs, - extra_config=extra_config, - ) - context_params = outputs["context"].disaggregated_params - assert context_params is not None - assert context_params.request_type == "context_only" - assert len(outputs["context"].token_ids) == 1 - assert context_params.ctx_request_id is not None - assert context_params.first_gen_tokens is not None - assert has_handoff_transport_metadata(context_params) - assert outputs["generation"].token_ids - assert has_draft_tokens(outputs["context"]) - assert has_draft_tokens(outputs["generation"]) - assert outputs["context"].text == " Berlin" - assert outputs["context"].token_ids == LLAMA_EAGLE3_EXPECTED_TOKEN_IDS[:1] - assert outputs["generation"].text == LLAMA_EAGLE3_EXPECTED_TEXT - assert outputs["generation"].token_ids == LLAMA_EAGLE3_EXPECTED_TOKEN_IDS diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_llama31_8b.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_llama31_8b.yaml deleted file mode 100644 index 487abd0d600e..000000000000 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_llama31_8b.yaml +++ /dev/null @@ -1,18 +0,0 @@ -hostname: localhost -model: llama-3.1-model/Llama-3.1-8B-Instruct -free_gpu_memory_fraction: 0.25 -backend: pytorch -disable_overlap_scheduler: true -gather_generation_logits: true -context_servers: - num_instances: 1 - tensor_parallel_size: 2 - pipeline_parallel_size: 1 - cache_transceiver_config: - backend: DEFAULT -generation_servers: - num_instances: 1 - tensor_parallel_size: 2 - pipeline_parallel_size: 1 - cache_transceiver_config: - backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 049cf3e3fc74..f521612400fa 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -28,7 +28,6 @@ from typing import Any, Optional import aiohttp -import numpy as np import pytest import yaml from defs.common import get_free_port_in_ci as get_free_port @@ -420,8 +419,6 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_cancel_stress_test.yaml", "cancel_stress_test_large": f"{test_configs_root}/disagg_config_cancel_stress_test_large.yaml", - "llama31_8b": - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_llama31_8b.yaml", "mamba_conc_greater_than_mbs": f"{test_configs_root}/disagg_config_mamba_conc_greater_than_mbs.yaml", "mamba_bs1_concurrency2": @@ -2352,9 +2349,6 @@ def benchmark_model_root(request): model_path = os.path.join(models_root, "DeepSeek-V3-Lite", "fp8") elif (request.param == "DeepSeek-V3-Lite-bf16"): model_path = os.path.join(models_root, "DeepSeek-V3-Lite", "bf16") - elif request.param == "llama-3.1-8b-instruct-hf-fp8": - model_path = os.path.join(models_root, "llama-3.1-model", - "Llama-3.1-8B-Instruct-FP8") else: raise ValueError(f"Failed to find the model: {request.param}") return model_path @@ -4075,205 +4069,6 @@ def test_disaggregated_cancel_large_context_requests(disaggregated_test_root, cwd=llm_venv.get_working_directory()) -@pytest.mark.skip_less_device(4) -@pytest.mark.parametrize("llama_model_root", ['llama-3.1-8b-instruct'], - indirect=True) -def test_disaggregated_logprobs_serving(disaggregated_test_root, - disaggregated_example_root, llm_venv, - llama_model_root): - """Test logprobs via OpenAI API in disaggregated serving with multi-GPU TP. - - Covers the RCCA scenario (NVBug 5926823): disaggregated + streaming + logprobs, - where the context worker returns prefill result (request_type=generation_only) - to the generation worker. Ensures LogProbStorage flows correctly across the - context/gen boundary without AttributeError on cum_log_probs. - """ - - async def iter_sse_chunks(resp): - """Yield parsed JSON chunks from an OpenAI SSE stream.""" - async for line in resp.content: - decoded = line.decode("utf-8").strip() - if not decoded.startswith("data: "): - continue - data_str = decoded[len("data: "):] - if data_str == "[DONE]": - break - try: - yield json.loads(data_str) - except json.JSONDecodeError: - continue - - async def collect_streaming_logprobs(resp, api_type): - """Parse SSE stream and return (tokens, logprobs) lists.""" - tokens, logprobs = [], [] - async for chunk in iter_sse_chunks(resp): - choices = chunk.get("choices", []) - if not choices: - continue - lp_data = choices[0].get("logprobs") - if not lp_data: - continue - if api_type == "completions": - tokens.extend(lp_data.get("tokens", [])) - logprobs.extend(lp_data.get("token_logprobs", [])) - else: - for item in lp_data.get("content", []): - tokens.append(item.get("token")) - logprobs.append(item.get("logprob")) - return tokens, logprobs - - def extract_logprobs(result, api_type): - """Extract (tokens, logprobs) from non-streaming OpenAI response.""" - choices = result.get("choices", []) - assert len(choices) > 0, "Response should have choices" - if api_type == "completions": - lp_data = choices[0].get("logprobs") - assert lp_data is not None, "Response should contain logprobs" - tokens = lp_data.get("tokens", []) - logprobs = lp_data.get("token_logprobs", []) - assert len(tokens) == len(logprobs), ( - f"count mismatch: {len(logprobs)} logprobs " - f"for {len(tokens)} tokens") - return tokens, logprobs - lp_obj = choices[0].get("logprobs") - assert lp_obj is not None, "Response should contain logprobs" - content = lp_obj.get("content", []) - tokens = [item.get("token") for item in content] - logprobs = [item.get("logprob") for item in content] - return tokens, logprobs - - setup_model_symlink(llm_venv, llama_model_root, - "llama-3.1-model/Llama-3.1-8B-Instruct") - - config_file = get_test_config("llama31_8b", disaggregated_example_root, - os.path.dirname(__file__)) - - env = llm_venv._new_env.copy() - env["UCX_TLS"] = get_ucx_tls() - ctx_workers, gen_workers, disagg_server, work_dir = [], [], None, None - config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ - setup_disagg_cluster(config_file, env=env, - model_name=llama_model_root, - cwd=llm_venv.get_working_directory(), - server_start_timeout=600) - - server_host = config.get("hostname", "localhost") - server_url = f"http://{server_host}:{server_port}" - model_name = "llama-3.1-model/Llama-3.1-8B-Instruct" - max_tokens = 20 - timeout = aiohttp.ClientTimeout(total=120) - # Use emoji prompt to also stress-test multi-byte tokenizer handling - prompt = "I love coding šŸš€ and AI." - - async def check_logprobs(): - async with aiohttp.ClientSession() as session: - for api_type in ("completions", "chat"): - url = (f"{server_url}/v1/completions" - if api_type == "completions" else - f"{server_url}/v1/chat/completions") - - def make_payload(prompt, stream, _api_type=api_type): - base = { - "max_tokens": max_tokens, - "logprobs": 1 if _api_type == "completions" else True, - "stream": stream, - "temperature": 0 - } - if _api_type == "completions": - return {"model": model_name, "prompt": prompt, **base} - return { - "model": model_name, - "messages": [{ - "role": "user", - "content": prompt - }], - **base - } - - # 1) Streaming vs non-streaming consistency check - async with session.post(url, - json=make_payload(prompt, False), - timeout=timeout) as resp: - assert resp.status == 200, \ - f"[{api_type}] non-streaming: {await resp.text()}" - ns_tokens, ns_logprobs = extract_logprobs( - await resp.json(), api_type) - - async with session.post(url, - json=make_payload(prompt, True), - timeout=timeout) as resp: - assert resp.status == 200, \ - f"[{api_type}] streaming: {await resp.text()}" - st_tokens, st_logprobs = \ - await collect_streaming_logprobs(resp, api_type) - - assert ns_tokens == st_tokens, ( - f"[{api_type}] streaming vs non-streaming tokens mismatch") - assert len(ns_logprobs) == len(st_logprobs), ( - f"[{api_type}] logprobs length: " - f"{len(ns_logprobs)} vs {len(st_logprobs)}") - # Skip position 0: the first token logprob can diverge - # between streaming and non-streaming in disaggregated mode - # due to the context/generation handoff boundary. - comparable = 0 - for i, (n, s) in enumerate( - zip(ns_logprobs, st_logprobs, strict=True)): - if i == 0 or n is None or s is None: - continue - comparable += 1 - rtol, atol = (1e-3, 1e-4) if api_type == "chat" else (1e-4, - 1e-5) - assert np.isclose(n, s, rtol=rtol, atol=atol), \ - f"[{api_type}] logprob mismatch at {i}: {n} vs {s}" - assert comparable > 0, ( - f"[{api_type}] no comparable post-handoff logprobs found") - - # 2) Chat API with top_logprobs (requires gather_generation_logits) - if api_type == "chat": - top_lp_payload = { - "model": model_name, - "messages": [{ - "role": "user", - "content": prompt - }], - "max_tokens": max_tokens, - "logprobs": True, - "top_logprobs": 3, - "stream": False, - "temperature": 0, - } - async with session.post(f"{server_url}/v1/chat/completions", - json=top_lp_payload, - timeout=timeout) as resp: - assert resp.status == 200, ( - f"[chat/top_logprobs] {resp.status}: " - f"{await resp.text()}") - result = await resp.json() - lp_obj = result["choices"][0].get("logprobs") - assert lp_obj is not None, "top_logprobs response should have logprobs" - content = lp_obj.get("content", []) - assert len( - content) > 0, "top_logprobs content should be non-empty" - for item in content: - top_lps = item.get("top_logprobs") - assert top_lps is not None and len(top_lps) > 0, ( - f"top_logprobs should be non-empty when requested: {item}" - ) - for tl in top_lps: - assert "token" in tl and "logprob" in tl, ( - f"top_logprob entry missing token/logprob: {tl}" - ) - assert tl["logprob"] <= 0.0, ( - f"top_logprob {tl['logprob']} should be <= 0") - - try: - asyncio.run(check_logprobs()) - finally: - terminate(*ctx_workers, *gen_workers, disagg_server) - if work_dir: - shutil.rmtree(work_dir, ignore_errors=True) - - @pytest.mark.skip_less_device(8) @skip_pre_blackwell @pytest.mark.parametrize("model_path", ['DeepSeek-V3-0324-FP4']) diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index beec99cff550..59805c24e67c 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py @@ -14,7 +14,6 @@ from tensorrt_llm._utils import set_mpi_comm from tensorrt_llm.llmapi import (CacheTransceiverConfig, CudaGraphConfig, KvCacheConfig, MpiCommSession) -from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig def get_ucx_tls(): @@ -49,8 +48,6 @@ def get_ucx_tls(): MODEL_PATHS = { "DeepSeek-V3-Lite-fp8": "DeepSeek-V3-Lite/fp8", "TinyLlama-1.1B-Chat-v1.0": "llama-models-v2/TinyLlama-1.1B-Chat-v1.0", - "Llama-3.1-8B-Instruct": "llama-3.1-model/Llama-3.1-8B-Instruct/", - "EAGLE3-LLaMA3.1-Instruct-8B": "EAGLE3-LLaMA3.1-Instruct-8B", "Qwen3-8B-FP8": "Qwen3/Qwen3-8B-FP8", } @@ -510,118 +507,6 @@ def test_disaggregated_llama_context_capacity(model, enable_cuda_graph, print("All workers terminated.") -@pytest.mark.parametrize("model", ["Llama-3.1-8B-Instruct"]) -@pytest.mark.parametrize("spec_dec_model_path", ["EAGLE3-LLaMA3.1-Instruct-8B"]) -@pytest.mark.parametrize("generation_overlap", [False]) -@pytest.mark.parametrize("eagle3_one_model", [True, False]) -def test_disaggregated_spec_dec_batch_slot_limit(model, spec_dec_model_path, - generation_overlap, - eagle3_one_model): - # Test whether the batch slots are properly released when using speculative decoding - # with disaggregated serving. - spec_dec_config = Eagle3DecodingConfig( - speculative_model=model_path(spec_dec_model_path), - eagle3_one_model=eagle3_one_model, - max_draft_len=3) - - worker_pytorch_configs = [] - - # Context worker - worker_pytorch_configs.append( - dict(disable_overlap_scheduler=True, - speculative_config=spec_dec_config, - max_batch_size=1)) - - # Generation worker - worker_pytorch_configs.append( - dict(disable_overlap_scheduler=not generation_overlap, - speculative_config=spec_dec_config, - max_batch_size=1)) - - kv_cache_configs = [ - KvCacheConfig(max_tokens=128, - enable_block_reuse=False, - free_gpu_memory_fraction=0.4) for _ in range(2) - ] - cache_transceiver_configs = [ - CacheTransceiverConfig(backend="DEFAULT") for _ in range(2) - ] - model_names = [model_path(model) for _ in range(2)] - ranks = [0, 1] - worker_args = list( - zip(kv_cache_configs, cache_transceiver_configs, worker_pytorch_configs, - model_names, ranks)) - - port_name = mpi_publish_name() - - prompt = "What is the capital of Germany?" - mpi_info = MPI.Info.Create() - mpi_info.Set("oversubscribe", "true") - with MPIPoolExecutor(max_workers=2, - env={ - "UCX_TLS": get_ucx_tls(), - "UCX_MM_ERROR_HANDLING": "y", - "OMPI_MCA_rmaps_base_oversubscribe": "1" - }, - mpi_info=mpi_info) as executor: - futures = [] - try: - for worker_arg in worker_args: - future = executor.submit(worker_entry_point, *worker_arg) - futures.append(future) - except Exception as e: - print(f"Error in worker {worker_arg}: {e}") - raise e - - intercomm = None - try: - print("Launched all the workers.") - intercomm = mpi_initialize_intercomm(port_name) - - for _ in range(2): - intercomm.recv(tag=MPI_READY) - print("Received ready signal.") - max_tokens = 25 - - requests = [] - for _ in range(10): - requests.append( - (prompt, SamplingParams(max_tokens=1, ignore_eos=True), - DisaggregatedParams(request_type="context_only"))) - - intercomm.send(requests, dest=0, tag=MPI_REQUEST) - - for _ in range(len(requests)): - output = intercomm.recv(source=0, tag=MPI_RESULT) - assert output[0].disaggregated_params is not None - assert output[ - 0].disaggregated_params.request_type == "context_only" - assert len(output[0].token_ids) == 1 - - generation_request_disagg_params = output[ - 0].disaggregated_params - generation_request_disagg_params.request_type = "generation_only" - requests = [] - requests.append((prompt, - SamplingParams(max_tokens=max_tokens, - ignore_eos=True), - generation_request_disagg_params)) - - intercomm.send(requests, dest=1, tag=MPI_REQUEST) - output = intercomm.recv(source=1, tag=MPI_RESULT) - - except MPI.Exception as e: - print(f"MPI Error") - raise e - finally: - mpi_send_termination_request(intercomm) - - # Wait for all futures to complete - for future in futures: - future.result() - print("All workers terminated.") - - @pytest.mark.parametrize("model", ["TinyLlama-1.1B-Chat-v1.0"]) @pytest.mark.parametrize("generation_overlap", [False, True]) def test_disaggregated_logprobs(model, generation_overlap): diff --git a/tests/integration/defs/examples/test_ad_speculative_decoding.py b/tests/integration/defs/examples/test_ad_speculative_decoding.py index 5a34bd58d09f..37ef41ef9e32 100644 --- a/tests/integration/defs/examples/test_ad_speculative_decoding.py +++ b/tests/integration/defs/examples/test_ad_speculative_decoding.py @@ -14,1087 +14,14 @@ # limitations under the License. import json -import os import re -from dataclasses import dataclass from pathlib import Path -from typing import Optional, Set -import pytest import torch -import torch.nn as nn -from defs.conftest import llm_models_root from test_common.llm_data import hf_id_to_local_model_dir -from transformers import AutoModelForCausalLM, AutoTokenizer -from transformers.masking_utils import create_causal_mask -from transformers.modeling_outputs import BaseModelOutputWithPast -from transformers.models.llama.modeling_llama import LlamaModel -from transformers.utils.generic import ModelOutput -from tensorrt_llm import SamplingParams -from tensorrt_llm._torch.auto_deploy.llm import LLM -from tensorrt_llm._torch.auto_deploy.models.custom.modeling_eagle import ( - EagleDrafterForCausalLM, - EagleWrapper, - EagleWrapperConfig, -) +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_eagle import EagleDrafterForCausalLM from tensorrt_llm._torch.auto_deploy.models.eagle import EagleDrafterFactory -from tensorrt_llm.llmapi import Eagle3DecodingConfig - -prompts = [ - "What is the capital of France?", - "Please explain the concept of gravity in simple words and a single sentence.", - "What are the main differences between Python and C++?", - "Summarize the plot of Romeo and Juliet in three sentences.", -] - -EAGLE_MODEL_SUBPATH = "EAGLE3-LLaMA3.1-Instruct-8B" -LLAMA_BASE_SUBPATH = "llama-3.1-model/Llama-3.1-8B-Instruct" -EAGLE_MAX_DRAFT_LEN = 3 - - -def get_model_paths(): - """Get model paths using llm_models_root().""" - models_root = llm_models_root() - base_model = os.path.join(models_root, LLAMA_BASE_SUBPATH) - eagle_model = os.path.join(models_root, EAGLE_MODEL_SUBPATH) - - print(f"Base model path: {base_model}") - print(f"EAGLE model path: {eagle_model}") - return base_model, eagle_model - - -@pytest.mark.parametrize( - ("attn_backend", "compile_backend"), - [ - ("trtllm", "torch-cudagraph"), - ("flashinfer", "torch-simple"), - ], -) -def test_autodeploy_eagle3_one_model_acceptance_rate(attn_backend: str, compile_backend: str): - """Test Eagle3 one-model acceptance rate with AutoDeploy engine. - - Runs Eagle3 one-model speculative decoding with streaming and verifies - that the acceptance rate is above a minimum threshold. - Parameterized over attention backend and compile backend. - """ - print("\n" + "=" * 80) - print( - f"Testing AutoDeploy Eagle3 One-Model Acceptance Rate " - f"(attn_backend={attn_backend}, compile_backend={compile_backend})" - ) - print("=" * 80) - - base_model, eagle_model = get_model_paths() - - print(f"\nBase Model: {base_model}") - print(f"Eagle3 Model: {eagle_model}") - - max_draft_len = EAGLE_MAX_DRAFT_LEN - - speculative_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model, - eagle3_one_model=True, - eagle3_layers_to_capture={1, 15, 28}, - ) - - with LLM( - model=base_model, - skip_loading_weights=False, - runtime="trtllm", - world_size=1, - speculative_config=speculative_config, - # Force the Eagle3 draft to match the target (Llama 3.1 8B is bfloat16). - # Shared KV cache requires matching dtypes between target and draft. - speculative_model_kwargs={"torch_dtype": "bfloat16"}, - compile_backend=compile_backend, - attn_backend=attn_backend, - transforms={"compile_model": {"piecewise_enabled": False}}, - max_num_tokens=512, - # max_batch_size must leave room for an extend-only sample batch during - # resize_kv_cache, i.e. max_num_tokens // max_batch_size >= 1 + max_draft_len. - # Otherwise the sample batch is classified as decode-only and the Eagle - # wrapper rejects it ("decode without drafting is not supported"). - # TODO: remove once resize_kv_cache is spec-aware. - # See: https://github.com/NVIDIA/TensorRT-LLM/issues/13348 - max_batch_size=128, - ) as llm: - _run_acceptance_rate_check(llm, max_draft_len) - - -def _run_acceptance_rate_check(llm, max_draft_len: int, min_acceptance_rate: float = 0.10): - """Common helper for acceptance rate tests. - - Submits all requests simultaneously so the executor processes them concurrently - (batch size > 1), then consumes streaming results to compute acceptance rates. - """ - batch_tok_ids = [llm.tokenizer.encode(p) for p in prompts] - sampling_params = SamplingParams(max_tokens=128, temperature=0, seed=42) - - print("\nRunning Eagle3 speculative decoding with streaming...") - print(f"Submitting all {len(batch_tok_ids)} requests simultaneously...") - - # Submit all requests before consuming any results so they are in-flight concurrently. - generators = [ - llm.generate_async(tok_ids, sampling_params, streaming=True) for tok_ids in batch_tok_ids - ] - - for i, gen in enumerate(generators): - num_tokens = 0 - num_drafted = 0 - num_accepted = 0 - - for output in gen: - new_tokens = output.outputs[0].token_ids - num_drafted += max_draft_len - num_accepted += len(new_tokens) - num_tokens - 1 - num_tokens = len(new_tokens) - - accept_rate = num_accepted / num_drafted - - generated_text = output.outputs[0].text - if not generated_text: - generated_text = llm.tokenizer.decode(output.outputs[0].token_ids) - print(f"\n[PROMPT {i}] {prompts[i]}") - print(f"[OUTPUT {i}] {generated_text}") - - print(f"\nRequest {i + 1} Acceptance Rate Statistics:") - print(f" Total tokens drafted: {num_drafted}") - print(f" Total tokens accepted: {num_accepted}") - print(f" Acceptance rate: {accept_rate:.2%}") - - assert accept_rate > min_acceptance_rate, ( - f"Request {i + 1}: Acceptance rate {accept_rate:.2%} is below minimum threshold " - f"{min_acceptance_rate:.0%}" - ) - - print("\n" + "=" * 80) - print("SUCCESS! All requests passed acceptance rate threshold") - print("=" * 80) - - -def load_weights(model_path: Path, model: torch.nn.Module): - """Load weights from checkpoint while applying the same _checkpoint_conversion_mapping that the factory uses. - - Returns: tuple of (loaded_keys, missing_keys, unexpected_keys) - """ - # 1. Load checkpoint keys - bin_path = model_path / "pytorch_model.bin" - safetensors_path = model_path / "model.safetensors" - - if safetensors_path.exists(): - from safetensors import safe_open - - with safe_open(safetensors_path, framework="pt") as f: - checkpoint_keys_original = list(f.keys()) - elif bin_path.exists(): - state_dict = torch.load(bin_path, map_location="cpu", weights_only=True) - checkpoint_keys_original = list(state_dict.keys()) - del state_dict - else: - raise FileNotFoundError(f"No checkpoint found at {model_path}") - - # 2. Apply _checkpoint_conversion_mapping (same logic as hf.py _remap_param_names_load_hook) - # This is the key part - the factory does this exact same thing in lines 496-512 of hf.py - conversion_mapping = getattr(model, "_checkpoint_conversion_mapping", None) - checkpoint_keys_remapped = [] - - for key in checkpoint_keys_original: - new_key = key - if conversion_mapping: - for pattern, replacement in conversion_mapping.items(): - new_key = re.sub(pattern, replacement, new_key) - checkpoint_keys_remapped.append(new_key) - - # 3. Get model's expected keys - model_keys = set(model.state_dict().keys()) - checkpoint_keys = set(checkpoint_keys_remapped) - - # 4. Calculate differences - loaded_keys = checkpoint_keys & model_keys - missing_in_checkpoint = model_keys - checkpoint_keys - unexpected_in_checkpoint = checkpoint_keys - model_keys - - return loaded_keys, missing_in_checkpoint, unexpected_in_checkpoint - - -def test_eagle_model_with_weights(): - """Test EagleModel forward pass with loaded weights using the EagleDrafterFactory. - - This test uses EagleDrafterFactory to initialize the model, which directly - builds the Eagle drafter model based on the checkpoint's model_type: - - 1. Factory creates config via AutoConfig.from_pretrained - 2. Factory selects EagleDrafterForCausalLM based on model_type="llama" - 3. Factory creates model via _from_config - 4. Factory loads weights via load_or_random_init -> _load_checkpoint - - This ensures the test validates the exact initialization path used in production. - """ - print("\n" + "=" * 80) - print("Test: EagleModel forward pass with loaded weights (via EagleDrafterFactory)") - print("=" * 80) - - _, eagle_model_path = get_model_paths() - eagle_path = Path(eagle_model_path) - - # 1. Setup Device - device = "cuda" if torch.cuda.is_available() else "cpu" - - # 2. Create factory - # EagleDrafterFactory directly builds the correct drafter model based on model_type - print("Creating EagleDrafterFactory...") - factory = EagleDrafterFactory( - model=eagle_model_path, - skip_loading_weights=False, # We want to test weight loading - ) - - # 3. Build model using factory - # Factory flow: - # build_model() -> prefetch_checkpoint() -> _build_model() - # _build_model() -> _get_model_config() (gets base LlamaConfig) - # _build_model() -> selects EagleDrafterForCausalLM for model_type="llama" - # _build_model() -> EagleDrafterForCausalLM._from_config(config) - print("Building model via factory.build_model('meta')...") - model = factory.build_model("meta") - print(f"Model type: {type(model).__name__}") - print(f"Model config type: {type(model.config).__name__}") - - # 4. Load weights from checkpoint and compare to model's expected keys - print("\n--- Weight Loading Analysis ---") - loaded_keys, missing_keys, unexpected_keys = load_weights(eagle_path, model) - - print(f"Total model parameters: {len(loaded_keys) + len(missing_keys)}") - print(f"Total checkpoint keys: {len(loaded_keys) + len(unexpected_keys)}") - print(f"āœ… Weights to be loaded: {len(loaded_keys)}") - print(f"āš ļø Missing in checkpoint (will be random init): {len(missing_keys)}") - print(f"āš ļø Unexpected in checkpoint (will be ignored): {len(unexpected_keys)}") - - if unexpected_keys: - print("\nUnexpected keys (in checkpoint but model doesn't expect):") - for key in sorted(unexpected_keys): - if "t2d" in key: - print(f" - {key} (expected: not used in Eagle3 for Llama3.1-8B-Instruct)") - else: - print(f" - {key}") - - if loaded_keys: - print(f"\nLoaded keys ({len(loaded_keys)} total):") - for key in sorted(loaded_keys)[:20]: - print(f" - {key}") - if len(loaded_keys) > 20: - print(f" ... and {len(loaded_keys) - 20} more") - - print("--- End Weight Analysis ---\n") - - # Verify expected missing and unexpected keys - # These are the keys we expect based on Eagle3 architecture: - # - embed_tokens: shared from target model (not in Eagle checkpoint) - # - t2d: target-to-draft mapping, not used in Eagle3 (uses d2t instead) - expected_unexpected_keys = {"model.t2d"} - - assert len(missing_keys) == 0, ( - f"Expect all keys to be loaded.\nKeys that are missing: {missing_keys}\n" - ) - - assert unexpected_keys == expected_unexpected_keys, ( - f"Unexpected keys in checkpoint.\n" - f"Expected: {expected_unexpected_keys}\n" - f"Got: {unexpected_keys}\n" - f"Extra unexpected: {unexpected_keys - expected_unexpected_keys}\n" - f"Not unexpected (but expected): {expected_unexpected_keys - unexpected_keys}" - ) - - print("āœ… Weight loading analysis matches expected missing/unexpected keys!") - - # 5. Load weights using factory (mimics actual pipeline) - # If tensor shapes do not match with how they are used in the forward() function, we will - # get an error. - print("Loading weights via factory.load_or_random_init()...") - factory.load_or_random_init(model, device) - print("Weights loaded successfully via factory interface!") - - model.eval() - - -############################################################################### -# Set up to test the prefill-only version of the EagleWrapper model in test_eagle_wrapper_forward(). -# This helps us guarantee that the EagleWrapper model, before it enters AutoDeploy, is working correctly, -# The test does not rely on any TRTLLM logic. -############################################################################### -class PrefillOnlyEagleResourceManager: - """Simple resource manager for Eagle speculative decoding (prefill-only variant). - - Stores hidden states for use by draft loop in EagleWrapper.forward(). - """ - - def __init__( - self, - hidden_size: int, - num_capture_layers: int, - max_batch_size: int, - max_seq_len: int, - max_draft_len: int, - target_dtype: torch.dtype, - ): - # Buffer for hidden states from target model: [max_tokens, hidden_size * num_capture_layers] - # Uses the same flattened 2D format as hidden_states_cache_* runtime buffers. - self.hidden_states = torch.empty( - max_batch_size * (max_seq_len + max_draft_len), - hidden_size * num_capture_layers, - device="cuda", - dtype=target_dtype, - ) - - -class LlamaModelWithCapture(LlamaModel): - """LlamaModel that captures un-normalized hidden states from specified layers. - - Overwrites the base model's forward method to capture hidden states from specified layers. - Base model's forward method is otherwise copied from LlamaModel in HuggingFace. - Takes PrefillOnlyEagleResourceManager as an argument to store captured hidden states. - """ - - def __init__( - self, - config, - layers_to_capture: Optional[Set[int]] = None, - resource_manager: Optional[PrefillOnlyEagleResourceManager] = None, - ): - super().__init__(config) - # layers_to_capture: set of layer indices (0-indexed) to capture - # If None, capture all layers - if layers_to_capture is None: - self.layers_to_capture = set(range(config.num_hidden_layers)) - else: - self.layers_to_capture = set(layers_to_capture) - - self.resource_manager = resource_manager - - # Validate layer indices - for idx in self.layers_to_capture: - if idx < 0 or idx >= config.num_hidden_layers: - raise ValueError( - f"Layer index {idx} out of range. " - f"Model has {config.num_hidden_layers} layers (0 to {config.num_hidden_layers - 1})" - ) - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - cache_position: Optional[torch.LongTensor] = None, - **kwargs, - ) -> BaseModelOutputWithPast: - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - if cache_position is None: - # prefill only - no past key values. - cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device) - - if position_ids is None: - position_ids = cache_position.unsqueeze(0) - - causal_mask = create_causal_mask( - config=self.config, - input_embeds=inputs_embeds, - attention_mask=attention_mask, - cache_position=cache_position, - past_key_values=None, - position_ids=position_ids, - ) - - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - # Buffer to collect captured hidden states - captured_hidden_states = [] - - for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask, - position_ids=position_ids, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - - # Capture this layer's output if it's in our list - if layer_idx in self.layers_to_capture: - captured_hidden_states.append(hidden_states) - - # Apply final normalization for last_hidden_state - last_hidden_state = self.norm(hidden_states) - - # Store captured hidden states in resource manager if available - # Resource manager uses 2D flattened format: [max_tokens, hidden_size * num_capture_layers] - if self.resource_manager is not None and captured_hidden_states: - concatenated = torch.cat(captured_hidden_states, dim=-1) - batch_size, seq_len, total_hidden_size = concatenated.shape - assert self.resource_manager.hidden_states.shape[-1] == total_hidden_size, ( - f"Resource manager buffer last dim {self.resource_manager.hidden_states.shape[-1]} " - f"!= concatenated hidden states last dim {total_hidden_size}" - ) - # Flatten to [batch_size * seq_len, total_hidden_size] for 2D format - flattened = concatenated.view(batch_size * seq_len, total_hidden_size) - self.resource_manager.hidden_states[: (batch_size * seq_len), :].copy_(flattened) - - return BaseModelOutputWithPast( - last_hidden_state=last_hidden_state, - hidden_states=tuple(captured_hidden_states) if captured_hidden_states else None, - ) - - -@dataclass -class LlamaForCausalLMOutput(ModelOutput): - logits: Optional[torch.FloatTensor] = None - last_hidden_state: Optional[torch.FloatTensor] = None - hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None - - -class LlamaForCausalLMWithCapture(nn.Module): - """Wrapper combining LlamaModelWithCapture with lm_head for EagleWrapper testing. - - EagleWrapper.forward() expects target_model(input_ids, position_ids) to return logits. - This class wraps LlamaModelWithCapture (which captures hidden states to resource manager) - and adds the lm_head to produce logits. - """ - - def __init__(self, base_model, capture_model): - super().__init__() - self.model = capture_model # LlamaModelWithCapture with resource_manager - self.lm_head = base_model.lm_head - - def forward( - self, - input_ids: Optional[torch.LongTensor] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - position_ids: Optional[torch.LongTensor] = None, - **kwargs, - ): - outputs = self.model( - input_ids=input_ids, inputs_embeds=inputs_embeds, position_ids=position_ids, **kwargs - ) - logits = self.lm_head(outputs.last_hidden_state) - return LlamaForCausalLMOutput( - logits=logits, - last_hidden_state=outputs.last_hidden_state, - hidden_states=outputs.hidden_states, - ) - - def get_input_embeddings(self): - return self.model.embed_tokens - - def get_output_embeddings(self): - return self.model.lm_head - - @classmethod - def from_pretrained( - cls, - model_name: str, - resource_manager, - capture_layers, - dtype=torch.bfloat16, - ): - """Load a base model and create a LlamaForCausalLMWithCapture with shared weights.""" - print(f"Loading {model_name}...") - base_model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=dtype, - device_map={"": 0}, - ) - base_model.eval() - - # Create LlamaModelWithCapture that shares weights with the base model - original_llama_model = base_model.model - - capture_model = LlamaModelWithCapture.__new__(LlamaModelWithCapture) - nn.Module.__init__(capture_model) - - capture_model.config = original_llama_model.config - capture_model.layers_to_capture = capture_layers - capture_model.resource_manager = resource_manager - - # Share all modules (no weight copying) - capture_model.embed_tokens = original_llama_model.embed_tokens - capture_model.layers = original_llama_model.layers - capture_model.norm = original_llama_model.norm - capture_model.rotary_emb = original_llama_model.rotary_emb - capture_model.gradient_checkpointing = original_llama_model.gradient_checkpointing - - return cls(base_model, capture_model) - - -def build_eagle_wrapper( - base_model_path: str, - eagle_model_path: str, - resource_manager: PrefillOnlyEagleResourceManager, - capture_layers: Set[int], - max_seq_len: int, - max_draft_len: int, - target_dtype: torch.dtype, - device: torch.device, -) -> tuple[EagleWrapper, nn.Module]: - """Build an EagleWrapper model for testing. - - This function encapsulates the model building logic using manual model building. - - Returns: - A tuple of (eagle_wrapper, target_model) where: - - eagle_wrapper: The EagleWrapper model ready for inference. - - target_model: The target model (for verification steps). - """ - # Build EagleWrapper manually. - print("\n" + "-" * 40) - print("Building EagleWrapper") - print("-" * 40) - - # Create target model with capture - target_model = LlamaForCausalLMWithCapture.from_pretrained( - base_model_path, resource_manager, capture_layers, target_dtype - ) - print("āœ“ Created target model with capture") - - # Create draft model using EagleDrafterFactory (mimics production pipeline) - # This ensures weights are loaded correctly via the same path as AutoDeploy - print("\nCreating draft model via EagleDrafterFactory...") - draft_factory = EagleDrafterFactory( - model=eagle_model_path, - skip_loading_weights=False, - ) - - # Build model on meta device first, then load weights - draft_model = draft_factory.build_model("meta") - print(f" Model type: {type(draft_model).__name__}") - - # Load weights via factory - print(" Loading weights via factory.load_or_random_init()...") - draft_factory.load_or_random_init(draft_model, device) - draft_model.eval() - - # Create EagleWrapper config - wrapper_config = EagleWrapperConfig( - max_draft_len=max_draft_len, - load_embedding_from_target=draft_model.load_embedding_from_target, - load_lm_head_from_target=draft_model.load_lm_head_from_target, - ) - - # Build EagleWrapper (this also loads weights from target into draft model where necessary) - eagle_wrapper = EagleWrapper( - config=wrapper_config, - target_model=target_model, - draft_model=draft_model, - resource_manager=resource_manager, - ) - eagle_wrapper.eval() - print("āœ“ Built EagleWrapper") - - return eagle_wrapper, target_model - - -def generate_target_outputs( - target_model: nn.Module, - input_ids: torch.Tensor, - num_iterations: int, -) -> torch.Tensor: - """Generate tokens from target model using greedy sampling. - - Runs target_model.forward() in a loop, taking the last logit from each output, - applying greedy sampling with torch.argmax, and appending to input_ids. - - Args: - target_model: Model that returns logits from forward(input_ids, position_ids). - input_ids: Initial input token ids of shape [batch_size, seq_len]. - num_iterations: Number of tokens to generate. - - Returns: - output_ids: Tensor of shape [batch_size, seq_len + num_iterations] containing - the original input_ids plus the generated tokens. - """ - device = input_ids.device - init_seq_len = input_ids.shape[1] - print(f"Initial sequence length: {init_seq_len}") - current_ids = input_ids.clone() - - with torch.no_grad(): - for _ in range(num_iterations): - # Generate position_ids from current sequence length - seq_len = current_ids.shape[1] - position_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) - position_ids = position_ids.expand(current_ids.shape[0], -1) - - # Forward pass - logits = target_model(current_ids, position_ids=position_ids).logits - - # Take the last logit and apply greedy sampling - last_logits = logits[:, -1, :] # [batch_size, vocab_size] - next_token = torch.argmax(last_logits, dim=-1, keepdim=True) # [batch_size, 1] - - # Append to current_ids - current_ids = torch.cat([current_ids, next_token], dim=1) - - return current_ids - - -def print_token_analysis( - input_ids: torch.Tensor, - num_previously_accepted: torch.Tensor, - target_output_ids: torch.Tensor, - tokenizer, -) -> None: - """Print debug analysis of accepted vs speculative tokens for each batch. - - Args: - input_ids: Current input token ids of shape [batch_size, seq_len]. - num_previously_accepted: Number of accepted tokens per batch [batch_size]. - target_output_ids: Reference output from target model [batch_size, total_seq_len]. - tokenizer: Tokenizer for decoding tokens to text. - """ - batch_size = input_ids.shape[0] - print("\n --- Token Analysis (per batch) ---") - - for i in range(batch_size): - prev_accepted_i = num_previously_accepted[i].item() - - # Accepted tokens (before speculation): input_ids[i, :num_previously_accepted[i]] - accepted_tokens = input_ids[i, :prev_accepted_i] - # Speculative tokens: input_ids[i, num_previously_accepted[i]:] - speculative_tokens = input_ids[i, prev_accepted_i:] - - # Target model's expected token at this position - if prev_accepted_i < target_output_ids.shape[1]: - target_token_at_pos = target_output_ids[i, prev_accepted_i] - else: - target_token_at_pos = None - - print(f"\n Batch {i}:") - print(f" num_previously_accepted: {prev_accepted_i}") - print( - f" Accepted tokens ({accepted_tokens.shape[0]} tokens): {accepted_tokens.tolist()}" - ) - accepted_text = tokenizer.decode(accepted_tokens, skip_special_tokens=True) - print(f' Accepted text: "{accepted_text}"') - print( - f" Speculative tokens ({speculative_tokens.shape[0]} tokens): {speculative_tokens.tolist()}" - ) - if speculative_tokens.shape[0] > 0: - spec_text = tokenizer.decode(speculative_tokens, skip_special_tokens=False) - print(f' Speculative text: "{spec_text}"') - if target_token_at_pos is not None: - target_tok_id = target_token_at_pos.item() - target_tok_str = tokenizer.decode([target_tok_id]) - print( - f' Target model\'s next token at pos {prev_accepted_i}: {target_tok_id} ("{target_tok_str}")' - ) - - -def manual_sample_and_verify( - next_target_inputs: list, - num_accepted_tokens: torch.Tensor, - target_model: nn.Module, - eagle_wrapper: nn.Module, - max_draft_len: int, - device: torch.device, -) -> list: - """Manually verify speculative tokens using sample_and_verify. - - This is used for batch_size > 1 where truncation prevents speculative tokens - from being fed back, so we verify them manually before truncation. - - Args: - next_target_inputs: List of tensors, one per batch element. - num_accepted_tokens: Number of tokens accepted so far per batch [batch_size]. - target_model: The target model for running forward pass. - eagle_wrapper: The EagleWrapper containing sample_and_verify. - max_draft_len: Maximum draft length (for capping counts). - device: Device to run on. - - Returns: - List of (num_accepted, num_speculative) tuples for each batch element. - """ - batch_size = len(next_target_inputs) - - # Due to our truncation trick, all sequences should have the same length - seq_lens = [seq.shape[0] for seq in next_target_inputs] - assert all(slen == seq_lens[0] for slen in seq_lens), ( - f"All sequences should have same length due to truncation, got {seq_lens}" - ) - verify_seq_len = seq_lens[0] - - # Stack into batched tensor - stacked_inputs = torch.stack(next_target_inputs, dim=0) # [batch_size, seq_len] - - # Run target model forward to get logits - verify_position_ids = ( - torch.arange(verify_seq_len, device=device, dtype=torch.long) - .unsqueeze(0) - .expand(batch_size, -1) - ) - with torch.no_grad(): - verify_target_logits = target_model(stacked_inputs, position_ids=verify_position_ids).logits - - # new_num_previously_accepted = num_accepted_tokens + 1 - # This represents the tokens accepted after target model's output from this iteration - new_num_previously_accepted = num_accepted_tokens + 1 - - # Call sample_and_verify to get acceptance counts - _, verify_newly_accepted, _, _ = eagle_wrapper.sample_and_verify( - stacked_inputs, verify_target_logits, new_num_previously_accepted - ) - - # Build results list - results = [] - for i in range(batch_size): - num_accepted_i = min(verify_newly_accepted[i].item(), max_draft_len) - num_speculative = next_target_inputs[i].shape[0] - new_num_previously_accepted[i].item() - results.append((num_accepted_i, num_speculative)) - - return results - - -def verify_eagle_wrapper_output(output, tokenizer, batch_size, num_previously_accepted): - """Verify the output structure and values from EagleWrapper forward pass. - - Args: - output: The output from EagleWrapper forward pass. - tokenizer: The tokenizer for decoding tokens. - batch_size: The batch size. - num_previously_accepted: Tensor of previously accepted token counts. - """ - # Verify output structure - print("\nOutput verification:") - assert output is not None, "Output should not be None" - assert hasattr(output, "new_tokens"), "Output should have new_tokens" - assert hasattr(output, "new_tokens_lens"), "Output should have new_tokens_lens" - - print(f" new_tokens: {type(output.new_tokens)} with {len(output.new_tokens)} items") - for i, tokens in enumerate(output.new_tokens): - new_tokens_text = tokenizer.decode(tokens, skip_special_tokens=True) - print(f" batch {i}: shape {tokens.shape}, tokens: {tokens.tolist()}") - print(f' batch {i}: decoded: "{new_tokens_text}"') - - # Compute num_accepted_tokens from new_tokens_lens + num_previously_accepted - num_accepted_tokens = num_previously_accepted + output.new_tokens_lens - - print(f" new_tokens_lens: {output.new_tokens_lens}") - print(f" num_accepted_tokens (computed): {num_accepted_tokens}") - - # Verify new_tokens_lens is within expected bounds - assert output.new_tokens_lens.shape == (batch_size,), ( - f"new_tokens_lens shape should be ({batch_size},), got {output.new_tokens_lens.shape}" - ) - - -@pytest.mark.skip( - reason="EagleWrapper interface was refactored (resource_manager removed from __init__, " - "sample_and_verify removed); test needs to be updated to match the new interface. " - "This test is valuable for validating Eagle3 correctness (acceptance ratio) directly " - "on the EagleWrapper model *before* the full export + transforms + KV-cache pipeline, " - "making it much easier to debug Eagle3 model issues in isolation. TODO: rewrite to " - "match the current EagleWrapper prefill-only and KV-cache forward interfaces." -) -@pytest.mark.parametrize("batch_size", [1, 2]) -def test_eagle_wrapper_forward(batch_size: int): - """Test EagleWrapper forward pass with target and draft models. - - This test validates the full speculative decoding loop: - 1. Target model processes input and captures hidden states - 2. Draft model generates speculative tokens - 3. EagleWrapper orchestrates verification and drafting - - For batch size 1, we call EagleWrapper forward in the expected way. Each iteration generates a "golden token" - (target output) and draft tokens. We input all of them to the wrapper model, - which verifies the draft tokens against the target output. It then outputs the accepted tokens - and newly generated draft tokens, along with numbers of accepted tokens, and the process repeats. - - For batch size > 1, we need to work around the fact that as we run the loop described above, the sequences lengths - in the batch will get out of sync. So instead, we do not provide validated draft tokens as input in each iteration - - we just input the first accepted token from the previous iteration - (which we know was generated by the target model), which keeps the batches in sync. - - To verify that the output draft tokens are reasonable, we run a manual target model verification step - after each iteration. We record how many of the output draft tokens were accepted. - - In the end, we test that the acceptance ratio of the draft tokens generated by the EagleWrapper is reasonable. - - Args: - batch_size: Number of prompts to process in parallel. - """ - print("\n" + "=" * 80) - print("Test: EagleWrapper forward pass") - print("=" * 80) - - # Set random seeds for reproducibility - torch.manual_seed(42) - - # Get model paths using integration test conventions - base_model_path, eagle_model_path = get_model_paths() - eagle_path = Path(eagle_model_path) - - # Configuration - capture_layers = {1, 15, 28} # Layers to capture for Eagle3 - num_capture_layers = len(capture_layers) - hidden_size = 4096 # Llama 3.1-8B hidden size - dtype = torch.bfloat16 - device = torch.device("cuda") - - # Test dimensions - max_batch_size = 4 - max_seq_len = 1024 - max_draft_len = 3 - - # Tokenize the test prompts - tokenizer = AutoTokenizer.from_pretrained(base_model_path) - # Llama uses left padding for batch inference - tokenizer.pad_token = tokenizer.eos_token - tokenizer.padding_side = "left" - - if batch_size == 1: - input_ids = tokenizer.encode(prompts[0], return_tensors="pt").to(device) - else: - tokenized = tokenizer( - prompts[:batch_size], - return_tensors="pt", - padding=True, - ) - input_ids = tokenized.input_ids.to(device) - - print(f"input_ids: {input_ids}") - seq_len = input_ids.shape[1] - init_seq_len = seq_len # Store initial sequence length for final comparison - - print("\nTest configuration:") - print(f" target_model: {base_model_path}") - print(f" eagle_model: {eagle_path}") - print(f" batch_size: {batch_size}, seq_len: {seq_len}") - print(f" max_draft_len: {max_draft_len}") - print(f" capture_layers: {capture_layers}") - print(f" prompts: {prompts[:batch_size]}") - print(f" input_ids: {input_ids}") - - # Create resource manager - resource_manager = PrefillOnlyEagleResourceManager( - hidden_size=hidden_size, - num_capture_layers=num_capture_layers, - max_batch_size=max_batch_size, - max_seq_len=max_seq_len, - max_draft_len=max_draft_len, - target_dtype=dtype, - ) - print("\nāœ“ Created resource manager") - print(f" target_hidden_states shape: {resource_manager.hidden_states.shape}") - - # Build eagle_wrapper and target_model using the refactored function - eagle_wrapper, target_model = build_eagle_wrapper( - base_model_path=base_model_path, - eagle_model_path=str(eagle_path), - resource_manager=resource_manager, - capture_layers=capture_layers, - max_seq_len=max_seq_len, - max_draft_len=max_draft_len, - target_dtype=dtype, - device=device, - ) - - # Create test inputs (input_ids already created from tokenizer above) - position_ids = ( - torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1) - ) - # Set previously_accepted_tokens to the length of input_ids (all context tokens are accepted) - # Shape should be [batch_size] - a 1D tensor with one value per batch - num_previously_accepted = torch.full((batch_size,), seq_len, device=device, dtype=torch.long) - - print("\nTest inputs:") - print(f" input_ids shape: {input_ids.shape}") - print(f" input_ids: {input_ids}") - print(f" position_ids shape: {position_ids.shape}") - print(f" num_previously_accepted: {num_previously_accepted}") - - # Generate target model outputs with greedy sampling - print("\nGenerating target model outputs with greedy sampling (for verification)...") - target_output_ids = generate_target_outputs(target_model, input_ids, num_iterations=100) - print(f" target_output_ids shape: {target_output_ids.shape}") - print(f" target_output_ids: {target_output_ids}") - - # Decode to text as sanity check - generated_text = tokenizer.decode(target_output_ids[0], skip_special_tokens=True) - print(f"\n Target model greedy generation decoded text:\n {generated_text}") - - print("\nāœ“ EagleWrapper forward pass completed successfully!") - print("āœ“ Output structure verified") - print("āœ“ new_tokens_lens within expected bounds") - print("āœ“ Target model greedy generation completed") - - print("\n================================================") - - num_iterations = 70 - - # Dictionary to track distribution of new_tokens_lens - # keys: 0 to max_draft_len - # newly_accepted_counts[i]: number of times the number of accepted draft tokens was i - newly_accepted_counts = {i: 0 for i in range(max_draft_len + 1)} - - for iteration in range(num_iterations): - print(f"\n{'=' * 40}") - print(f"EagleWrapper forward pass - Iteration {iteration + 1}/{num_iterations}") - print(f"{'=' * 40}") - - seq_len = input_ids.shape[1] - - # Debug: Print speculative tokens, accepted tokens, and target comparison - print_token_analysis(input_ids, num_previously_accepted, target_output_ids, tokenizer) - - kwargs = { - "num_previously_accepted": num_previously_accepted, - } - with torch.no_grad(): - output = eagle_wrapper( - input_ids=input_ids, - position_ids=position_ids, - **kwargs, - ) - - verify_eagle_wrapper_output(output, tokenizer, batch_size, num_previously_accepted) - - # Prepare next_target_inputs - # output.new_tokens[i] contains the full draft_input_ids tensor, but the valid prefix - # has length num_accepted_tokens[i] + max_draft_len. We slice to get only valid tokens. - # We then prepend the first token from the previous iteration's input_ids. - # This prepending is only needed for prefill-only mode, since in the cached case, the first token - # will always be in the KV cache. - # Compute num_accepted_tokens from num_previously_accepted + new_tokens_lens - num_accepted_tokens = num_previously_accepted + output.new_tokens_lens - valid_prefix_len = num_accepted_tokens + max_draft_len - next_target_inputs = [ - torch.cat( - (input_ids[i, 0].unsqueeze(0), output.new_tokens[i][: valid_prefix_len[i]]), - dim=0, - ) - for i in range(batch_size) - ] - - # Track distribution of newly accepted tokens by reading new_tokens_lens from the output. - # For batch size = 1, we are inputting draft tokens to the wrapper model, so new_tokens_lens - # gives the number of accepted tokens from drafts in the previous iteration. - if batch_size == 1: - for val in output.new_tokens_lens.tolist(): - newly_accepted_counts[val] += 1 - print(f" newly_accepted_counts so far: {newly_accepted_counts}") - - # For batch_size > 1, we use manual target model verification below instead to check which of the draft tokens - # generated in *this* iteration would be accepted by the target model. - else: - # For batch_size > 1, verify acceptance using sample_and_verify() - # before truncation (since truncation prevents speculative tokens from being fed back) - verify_results = manual_sample_and_verify( - next_target_inputs, - num_accepted_tokens, - target_model, - eagle_wrapper, - max_draft_len, - device, - ) - - # Update newly_accepted_counts map - for i, (num_accepted_i, num_speculative) in enumerate(verify_results): - newly_accepted_counts[num_accepted_i] += 1 - print( - f" [Batch {i}] sample_and_verify: {num_accepted_i}/{num_speculative} speculative accepted" - ) - - # Truncate to keep shapes consistent across batches in each iteration. - # We know that the first token that is generated in this iteration is accepted, so it is "safe". - # All speculative tokens are truncated regardless of whether they are accepted or not. - # This is a hack to prevent the sequence lengths from getting out of sync across batches in each iteration - # without needing to change the padding every iteration. - truncate_len = input_ids.shape[1] + 1 - next_target_inputs = [seq[:truncate_len] for seq in next_target_inputs] - - next_target_inputs = torch.stack(next_target_inputs, dim=0) - - print(f" next_target_inputs: {next_target_inputs}") - print(f" next_target_inputs.shape: {next_target_inputs.shape}") - - # Update for next iteration - input_ids = next_target_inputs - seq_len = input_ids.shape[1] - position_ids = torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0) - position_ids = position_ids.expand(batch_size, -1) - - if batch_size > 1: - # For multi-batch: increment by 1 (we truncated, so just advance by one token) - num_previously_accepted = num_previously_accepted + 1 - else: - # For single batch: accept the tokens accepted in the previous iteration, plus one - # for the output token that was generated by the target. - num_previously_accepted = num_accepted_tokens + 1 - - print(f"\n{'=' * 40}") - print(f"Loop completed: {num_iterations} iterations") - print("Newly accepted tokens distribution:") - for k, v in newly_accepted_counts.items(): - print(f" {k}: {v}") - - # Calculate acceptance ratio - # For batch_size == 1: uses new_tokens_lens from eagle wrapper - # For batch_size > 1: uses manual verification against target model (since truncation - # prevents speculative tokens from being fed back) - total_accepted = sum(k * v for k, v in newly_accepted_counts.items()) - # First iteration has no tokens to newly accept, subsequent iterations have max_draft_len potential - - num_iterations_with_drafts = num_iterations - 1 if batch_size == 1 else num_iterations - total_potential = max_draft_len * (num_iterations_with_drafts) * batch_size - acceptance_ratio = total_accepted / total_potential if total_potential > 0 else 0.0 - print(f"\nAcceptance ratio: {total_accepted}/{total_potential} = {acceptance_ratio:.3f}") - if batch_size > 1: - print(" (batch_size > 1: measured via manual target model verification)") - assert acceptance_ratio > 0.1, ( - f"Acceptance ratio {acceptance_ratio:.3f} is too low (expected > 0.1)" - ) - - print("\n" + "=" * 80) - print("FINAL OUTPUT COMPARISON") - print("=" * 80) - for i in range(batch_size): - print(f"\n{'ā”€' * 40}") - print(f"BATCH {i}") - print(f"{'ā”€' * 40}") - print(f"\n[Target Model Output] ({target_output_ids[i].shape[0]} tokens):") - print(f" Tokens: {target_output_ids[i].tolist()}") - print(f' Text: "{tokenizer.decode(target_output_ids[i], skip_special_tokens=True)}"') - print(f"\n[Eagle Wrapper Output] ({input_ids[i].shape[0]} tokens):") - print(f" Tokens: {input_ids[i].tolist()}") - print(f' Text: "{tokenizer.decode(input_ids[i], skip_special_tokens=True)}"') - print("\n" + "=" * 80) - - # Verify that the first 10 generated tokens match between target model and eagle wrapper - # They seem to diverge after awhile but are semantically the same. - # Note that even running the target model in decode vs prefill mode, the outputs seem to diverge similarly, - # so this is not worrisome. This test provides a check that they are "similar enough" to each other. - num_tokens_to_check = 10 - print(f"\nVerifying first {num_tokens_to_check} generated tokens match...") - for i in range(batch_size): - target_generated = target_output_ids[i, init_seq_len : init_seq_len + num_tokens_to_check] - eagle_generated = input_ids[i, init_seq_len : init_seq_len + num_tokens_to_check] - - print(f" Batch {i}:") - print(f" Target: {target_generated.tolist()}") - print(f" Eagle: {eagle_generated.tolist()}") - - assert torch.equal(target_generated, eagle_generated), ( - f"Batch {i}: First {num_tokens_to_check} generated tokens do not match!\n" - f" Target: {target_generated.tolist()}\n" - f" Eagle: {eagle_generated.tolist()}" - ) - print(f"āœ“ First {num_tokens_to_check} generated tokens match for all batches!") def _load_valid_safetensors_index(index_path: Path): diff --git a/tests/integration/defs/llmapi/test_llm_examples.py b/tests/integration/defs/llmapi/test_llm_examples.py index 36f12ccd1905..b30d022a1893 100644 --- a/tests/integration/defs/llmapi/test_llm_examples.py +++ b/tests/integration/defs/llmapi/test_llm_examples.py @@ -61,9 +61,6 @@ def _run_llmapi_example(llm_root, engine_dir, llm_venv, script_name: str, # medusa-vicuna-7b-v1.3 f"{llm_models_root()}/medusa-vicuna-7b-v1.3": f"{llm_venv.get_working_directory()}/FasterDecoding/medusa-vicuna-7b-v1.3", - # llama3.1-medusa-8b-hf_v0.1 - f"{llm_models_root()}/llama3.1-medusa-8b-hf_v0.1": - f"{llm_venv.get_working_directory()}/nvidia/Llama-3.1-8B-Medusa-FP8", # Llama-3.1-8B-Instruct f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct": f"{llm_venv.get_working_directory()}/meta-llama/Llama-3.1-8B-Instruct", diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 443a621f6742..7d32cfed017f 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -23,10 +23,10 @@ import pytest import yaml -from defs.trt_test_alternative import (check_call, check_call_negative_test, - check_output, print_info, print_warning) +from defs.trt_test_alternative import (check_call, check_output, print_info, + print_warning) -from .common import get_mmlu_accuracy, venv_check_call +from .common import get_mmlu_accuracy from .conftest import (get_device_count, get_sm_version, llm_models_root, skip_post_blackwell, skip_pre_ada, skip_pre_blackwell, skip_pre_hopper, unittest_path) @@ -383,178 +383,6 @@ def temp_extra_llm_api_options_file(request): yield None -@pytest.mark.parametrize( - "model_name, llama_model_root, use_extra_config, pytorch_backend_config", - [('meta-llama/Llama-3.1-8B', 'llama-3.1-8b', False, False), - pytest.param('meta-llama/Llama-3.1-8B', - 'llama-3.1-8b-instruct-hf-fp8', - True, - False, - marks=skip_pre_hopper), - pytest.param('meta-llama/Llama-3.1-8B', - 'llama-3.1-8b-instruct-hf-fp8', - True, - True, - marks=skip_pre_hopper), - pytest.param('meta-llama/Llama-3.1-8B', - 'llama-3.1-8b-hf-nvfp4', - False, - False, - marks=skip_pre_blackwell)], - indirect=['llama_model_root']) -def test_trtllm_bench_pytorch_backend_sanity(llm_root, llm_venv, - llama_model_root, model_name, - use_extra_config, - pytorch_backend_config, - temp_extra_llm_api_options_file): - """Sanity check on latency benchmark for LLM API with PyTorch backend - """ - model_path, dataset_path = trtllm_bench_prolog(llm_root, llm_venv, - llama_model_root, model_name, - False, False) - - benchmark_cmd = \ - f"trtllm-bench --model {model_name} --model_path {model_path} " \ - f"throughput " \ - f"--dataset {dataset_path} --backend pytorch" - - mapping = { - "Meta-Llama-3.1-8B": 19.4, - "Llama-3.1-8B-Instruct-FP8": 12.0, - "Meta-Llama-3.1-8B-NVFP4": 10.2 - } - if use_extra_config: - benchmark_cmd += f" --config {temp_extra_llm_api_options_file}" - - model_id = llama_model_root.split(r"/")[-1] - if "nvfp4-quantized" in llama_model_root: - model_id += "-NVFP4" - - check_call(benchmark_cmd, shell=True) - - -def test_trtllm_bench_mgmn(llm_root, llm_venv): - model_name = "meta-llama/Llama-3.1-8B" - llama_model_dir = Path( - llm_models_root()) / "llama-3.1-model/Llama-3.1-8B-Instruct" - _, dataset_path = trtllm_bench_prolog(llm_root, - llm_venv, - model_subdir=llama_model_dir, - model_name=model_name, - quant=None, - streaming=False) - - benchmark_cmd = \ - f"mpirun --allow-run-as-root -n 2 trtllm-llmapi-launch trtllm-bench --model {model_name} " \ - f"--model_path {llama_model_dir} " \ - f"throughput " \ - f"--dataset {str(dataset_path)} --backend pytorch --tp 2" - - check_call(benchmark_cmd, shell=True, env=llm_venv._new_env) - - -@pytest.mark.parametrize( - "model_name", - [ - "meta-llama/Llama-3.1-8B", - ], -) -def test_trtllm_bench_help_sanity(model_name): - """Sanity check that the options are defined properly by printing out help - """ - check_call("trtllm-bench --help", shell=True) - check_call(f"trtllm-bench --model {model_name} throughput --help", - shell=True) - check_call(f"trtllm-bench --model {model_name} latency --help", shell=True) - - -@pytest.mark.parametrize("request_rate", [False, True], - ids=["", "enable_request_rate"]) -@pytest.mark.parametrize("concurrency", [False, True], - ids=["", "enable_concurrency"]) -def test_trtllm_bench_request_rate_and_concurrency(llm_root, llm_venv, - request_rate, concurrency): - """Sanity check on the trtllm-bench new request rate and concurrency API - """ - model_subdir = "llama-3.1-model/Meta-Llama-3.1-8B" - model_name = "meta-llama/Llama-3.1-8B" - - model_path, dataset_path = trtllm_bench_prolog(llm_root, - llm_venv, - model_subdir, - model_name, - quant=None, - streaming=False) - - benchmark_cmd = \ - f"trtllm-bench --model {model_name} --model_path {model_path} throughput " \ - f"--dataset {dataset_path} --backend pytorch" - - if request_rate: - benchmark_cmd += " --request_rate 100" - if concurrency: - benchmark_cmd += " --concurrency 100" - - print(f"cmd: {benchmark_cmd}") - - if request_rate and concurrency: - # negative test, request rate and concurrency should not be turned on at the same time - check_call_negative_test(benchmark_cmd, shell=True) - else: - check_call(benchmark_cmd, shell=True) - - -@pytest.mark.parametrize("model_subdir", [ - "llama-3.1-model/Meta-Llama-3.1-8B", -], - ids=lambda x: x.strip("-")) -@pytest.mark.parametrize( - "model_name", - [ - "meta-llama/Llama-3.1-8B", - ], -) -@pytest.mark.parametrize("streaming", [True, False], - ids=["non-streaming", "streaming"]) -@pytest.mark.parametrize("backend", ["pytorch"], ids=["PyTorch"]) -def test_trtllm_bench_iteration_log(llm_root, llm_venv, model_name, - model_subdir, streaming, backend): - """Test the iteration log functionality with necessary options - """ - iteration_log = None - - try: - iteration_log = tempfile.mkstemp(dir="/tmp", suffix=".txt")[1] - - model_path, dataset_path = trtllm_bench_prolog(llm_root, - llm_venv, - model_subdir, - model_name, - quant=None, - streaming=streaming) - - benchmark_cmd = \ - f"trtllm-bench --model {model_name} --model_path {model_path} " \ - f"throughput --dataset {dataset_path} --iteration_log {iteration_log}" - - if streaming: - benchmark_cmd += " --streaming" - - benchmark_cmd += f" --backend {backend}" - - check_call(benchmark_cmd, shell=True) - - assert os.path.exists( - iteration_log - ), f"Iteration log file {iteration_log} was not created." - if os.path.getsize(iteration_log) == 0: - raise AssertionError( - f"Iteration log file {iteration_log} is empty.") - finally: - if iteration_log: - shutil.rmtree(iteration_log, ignore_errors=True) - - def test_trtllm_serve_example(llm_root, llm_venv): example_root = Path(os.path.join(llm_root, "examples", "serve")) test_root = unittest_path() / "llmapi" / "apps" @@ -791,10 +619,8 @@ def test_openai_mmencoder_example(llm_root, llm_venv): str(test_root / "_test_openai_mmencoder.py")]) -@pytest.mark.parametrize("model_name", [ - "meta-llama/Llama-3.1-8B-Instruct", - pytest.param("openai/gpt-oss-120b", marks=skip_pre_hopper) -]) +@pytest.mark.parametrize( + "model_name", [pytest.param("openai/gpt-oss-120b", marks=skip_pre_hopper)]) def test_openai_chat_guided_decoding(llm_root, llm_venv, model_name: str): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ @@ -805,10 +631,8 @@ def test_openai_chat_guided_decoding(llm_root, llm_venv, model_name: str): @pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize("model_name", [ - "llama-3.1-model/Meta-Llama-3.1-8B", - pytest.param("gpt_oss/gpt-oss-20b", marks=skip_pre_hopper) -]) +@pytest.mark.parametrize( + "model_name", [pytest.param("gpt_oss/gpt-oss-20b", marks=skip_pre_hopper)]) def test_trtllm_benchmark_serving(llm_venv, model_name): test_root = unittest_path() / "llmapi" / "apps" llm_venv.run_cmd([ @@ -828,39 +652,6 @@ def test_trtllm_multimodal_benchmark_serving(llm_root, llm_venv): ]) -@pytest.mark.skip_less_device(4) -@pytest.mark.skip_less_device_memory(40000) -@pytest.mark.parametrize("service_discovery", ["etcd"]) -def test_openai_disagg_multi_nodes_completion_service_discovery( - llm_root, llm_venv, service_discovery): - test_root = unittest_path() / "llmapi" / "apps" - llm_venv.run_cmd([ - "-m", - "pytest", - str(test_root / - f"_test_disagg_serving_multi_nodes_service_discovery.py::test_completion[{service_discovery}]" - ), - ]) - - -@pytest.mark.skip_less_device(4) -@pytest.mark.skip_less_device_memory(40000) -@pytest.mark.parametrize("gen_config", - ["gen_tp2pp1", "gen_tp1pp2", "gen_tp1pp1"]) -@pytest.mark.parametrize("ctx_config", - ["ctx_tp2pp1", "ctx_tp1pp2", "ctx_tp1pp1"]) -def test_openai_disagg_multi_nodes_completion(llm_root, llm_venv, ctx_config, - gen_config): - test_root = unittest_path() / "llmapi" / "apps" - llm_venv.run_cmd([ - "-m", - "pytest", - str(test_root / - f"_test_disagg_serving_multi_nodes.py::test_completion[{ctx_config}-{gen_config}]" - ), - ]) - - ### PyTorch examples @@ -880,27 +671,9 @@ def parse_output(text): return results -def test_ptp_quickstart(llm_root, llm_venv): - example_root = Path(os.path.join(llm_root, "examples", "llm-api")) - - src = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct" - dst = f"{llm_venv.get_working_directory()}/meta-llama/Llama-3.1-8B-Instruct" - os.makedirs(os.path.dirname(dst), exist_ok=True) - os.symlink(src, dst, target_is_directory=True) - - venv_check_call(llm_venv, [str(example_root / "quickstart_example.py")]) - - @pytest.mark.parametrize("model_name,model_path", [ - ("Llama3.1-8B-BF16", "llama-3.1-model/Meta-Llama-3.1-8B"), ("Llama3.2-11B-BF16", "llama-3.2-models/Llama-3.2-11B-Vision"), ("Nemotron4_4B-BF16", "nemotron/Minitron-4B-Base"), - pytest.param('Llama3.1-8B-NVFP4', - 'nvfp4-quantized/Meta-Llama-3.1-8B', - marks=skip_pre_blackwell), - pytest.param('Llama3.1-8B-FP8', - 'llama-3.1-model/Llama-3.1-8B-Instruct-FP8', - marks=skip_pre_hopper), pytest.param('Qwen3-30B-A3B', 'Qwen3/Qwen3-30B-A3B', marks=pytest.mark.skip_less_device_memory(80000)), @@ -916,10 +689,6 @@ def test_ptp_quickstart(llm_root, llm_venv): marks=skip_pre_blackwell), pytest.param( 'GPT-OSS-120B', 'gpt_oss/gpt-oss-120b', marks=skip_pre_blackwell), - ("Llama3.1-8B-bf16-instruct", "llama-3.1-model/Llama-3.1-8B-Instruct"), - pytest.param('Llama3.1-8B-FP4', - 'modelopt-hf-model-hub/Llama-3.1-8B-Instruct-fp4', - marks=skip_pre_blackwell), pytest.param( 'Qwen3-8b-fp8', 'Qwen3/nvidia-Qwen3-8B-FP8', marks=skip_pre_hopper), pytest.param('Qwen3-8b-nvfp4', @@ -958,11 +727,8 @@ def test_ptp_quickstart_advanced(llm_root, llm_venv, model_name, model_path): ]) else: mapping = { - "Llama3.1-8B-BF16": 18.60, "Llama3.2-11B-BF16": 18.88, "Nemotron4_4B-BF16": 12.50, - "Llama3.1-8B-FP8": 13.05, - "Llama3.1-8B-NVFP4": 10.2 } cmds = [ str(example_root / "quickstart_advanced.py"), @@ -1043,8 +809,6 @@ def test_ptp_quickstart_advanced_bs1(llm_root, llm_venv): @pytest.mark.parametrize("model_name,model_path,eagle_model_path", [ - ("Llama-3.1-8b-Instruct", "llama-3.1-model/Llama-3.1-8B-Instruct", - "EAGLE3-LLaMA3.1-Instruct-8B"), pytest.param('GPT-OSS-120B-Eagle3', 'gpt_oss/gpt-oss-120b', 'gpt_oss/gpt-oss-120b-Eagle3', @@ -1070,88 +834,6 @@ def test_ptp_quickstart_advanced_eagle3(llm_root, llm_venv, model_name, ]) -@pytest.mark.parametrize("model_name,model_path,eagle_model_path", [ - ("Llama-3.1-8b-Instruct", "llama-3.1-model/Llama-3.1-8B-Instruct", - "EAGLE3-LLaMA3.1-Instruct-8B"), -]) -def test_draft_token_tree_quickstart_advanced_eagle3(llm_root, llm_venv, - model_name, model_path, - eagle_model_path): - print(f"Testing {model_name}.") - example_root = Path(os.path.join(llm_root, "examples", "llm-api")) - llm_venv.run_cmd([ - str(example_root / "quickstart_advanced.py"), - "--prompt", - "You are a good assistant. Please tell me the capital of France is", - "--spec_decode_max_draft_len", - "3", - "--spec_decode_algo", - "eagle3", - "--model_dir", - f"{llm_models_root()}/{model_path}", - "--draft_model_dir", - f"{llm_models_root()}/{eagle_model_path}", - "--disable_kv_cache_reuse", - "--disable_overlap_scheduler", - "--eagle_choices", - "[[0], [1], [2], [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [2, 0], [0, 0, 0], [0, 1, 0], [1, 0, 0]]", - "--kv_cache_fraction", - "0.4", - ]) - - -@pytest.mark.parametrize("model_name,model_path,eagle_model_path", [ - ("Llama-3.1-8b-Instruct", "llama-3.1-model/Llama-3.1-8B-Instruct", - "EAGLE3-LLaMA3.1-Instruct-8B"), -]) -def test_draft_token_tree_quickstart_advanced_eagle3_depth_1_tree( - llm_root, llm_venv, model_name, model_path, eagle_model_path): - print(f"Testing {model_name}.") - example_root = Path(os.path.join(llm_root, "examples", "llm-api")) - llm_venv.run_cmd([ - str(example_root / "quickstart_advanced.py"), - "--prompt", - "You are a good assistant. Please tell me the capital of France is", - "--spec_decode_max_draft_len", - "3", - "--spec_decode_algo", - "eagle3", - "--model_dir", - f"{llm_models_root()}/{model_path}", - "--draft_model_dir", - f"{llm_models_root()}/{eagle_model_path}", - "--disable_kv_cache_reuse", - "--disable_overlap_scheduler", - "--eagle_choices", - "[[0], [1], [2]]", - "--kv_cache_fraction", - "0.4", - ]) - - -@pytest.mark.parametrize("model_name,model_path", [ - ("Llama-3.1-8B-Instruct", "llama-3.1-model/Llama-3.1-8B-Instruct"), -]) -def test_ptp_quickstart_advanced_ngram(llm_root, llm_venv, model_name, - model_path): - print(f"Testing {model_name}.") - example_root = Path(os.path.join(llm_root, "examples", "llm-api")) - llm_venv.run_cmd([ - str(example_root / "quickstart_advanced.py"), - "--model_dir", - f"{llm_models_root()}/{model_path}", - "--spec_decode_algo", - "NGRAM", - "--spec_decode_max_draft_len", - "4", - "--max_matching_ngram_size", - "2", - "--use_cuda_graph", - "--disable_kv_cache_reuse", - "--disable_overlap_scheduler", - ]) - - @skip_post_blackwell @pytest.mark.skip_less_device_memory(80000) @pytest.mark.skip_less_device(4) @@ -1418,23 +1100,6 @@ def test_ptp_quickstart_advanced_8gpus_chunked_prefill_sq_22k( llm_venv.run_cmd(cmd) -@skip_pre_blackwell -def test_ptp_quickstart_advanced_mixed_precision(llm_root, llm_venv): - example_root = Path(os.path.join(llm_root, "examples", "llm-api")) - model_path = "Llama-3_1-8B-Instruct_fp8_nvfp4_hf" - llm_venv.run_cmd([ - str(example_root / "quickstart_advanced.py"), - "--model_dir", - f"{llm_models_root()}/{model_path}", - ]) - - # NOTE: we deliberately do not check the LLM outputs with keyword matching ratios as in the - # other tests, as it can be brittle and cause flakiness in CI. - # This test now becomes a smoke / functional test. - # Proper accuracy tests should be added to - # `tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py`. - - @pytest.mark.parametrize("modality", ["image", "video"]) @pytest.mark.parametrize( "model_name,model_path,match_ratio", @@ -1868,58 +1533,3 @@ def test_get_ci_container_port(): assert container_port_start > 0 assert container_port_num > 0 assert container_port_start + container_port_num <= 60000 - - -@skip_pre_hopper -@pytest.mark.skip_less_device_memory(80000) -@pytest.mark.parametrize("model_name", ["meta/Meta-Llama-3.1-8B"], - ids=["llama3_1-8b"]) -@pytest.mark.parametrize("model_subdir", ["llama-3.1-model/Meta-Llama-3.1-8B"], - ids=["llama_v3_1"]) -def test_trtllm_bench_mig_launch(llm_root, llm_venv, model_name, model_subdir): - """Run benchmark in MIG mode, check if throughput increases with concurrency.""" - results = {} - concurrency_list = [1, 32, 64, 128] - - for concurrency in concurrency_list: - num_requests = concurrency * 10 - runner = BenchRunner(llm_root=llm_root, - llm_venv=llm_venv, - model_name=model_name, - model_subdir=model_subdir, - streaming=False, - use_mpirun=False, - tp_size=1, - concurrency=concurrency, - num_requests=num_requests) - - output = runner() - results[concurrency] = output - - print(f"\n=== Benchmark Results Comparison ===") - print(f"Model: {model_name}") - print( - f"{'Concurrency':<15} {'Throughput':<15} {'Latency':<15} {'Num Requests':<15}" - ) - print("-" * 60) - - for idx, val in enumerate(concurrency_list): - metrics = results.get(val) - if not isinstance(metrics, dict): - pytest.fail( - f"Unexpected benchmark result type for concurrency {val}: {type(metrics)}" - ) - try: - throughput = float(metrics.get('throughput', 0)) - latency = float(metrics.get('latency', 0)) - num_requests = int(metrics.get('num_requests', 0)) - except (ValueError, TypeError) as e: - pytest.fail( - f"Failed to parse benchmark results for concurrency {val}: {e}") - assert throughput > 0, f"Throughput is 0 for concurrency {val}" - assert latency > 0, f"Latency is 0 for concurrency {val}" - print(f"{val:<15} {throughput:<15} {latency:<15} {num_requests:<15}") - if idx > 0: - prev_throughput = float(results[concurrency_list[idx - 1]].get( - 'throughput', 0)) - assert throughput > prev_throughput * 1.3, f"Throughput is not increasing for concurrency {concurrency_list[idx]}" diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 2490f94f8613..6ba78745f57e 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -484,6 +484,8 @@ accuracy/test_llm_api_pytorch.py::TestLagunaXS_2_1::test_nvfp4_dflash accuracy/test_llm_api_pytorch.py::TestLlama4SpeculativeDecoding::test_guided_decoding TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestLlama4SpeculativeDecoding::test_llama4_eagle3[dynamic] TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestLlama4SpeculativeDecoding::test_llama4_eagle3[linear] TIMEOUT (60) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] +accuracy/test_llm_api_pytorch.py::TestMiniMaxM2_5::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] @@ -735,7 +737,6 @@ disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Ch disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] disaggregated/test_disaggregated.py::test_disaggregated_kv_cache_time_output[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0] -disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] disaggregated/test_disaggregated.py::test_disaggregated_mamba_bs1_concurrency2 disaggregated/test_disaggregated.py::test_disaggregated_mamba_conc_greater_than_mbs[NVIDIA-Nemotron-3-Super-120B-A12B-FP8] disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] @@ -775,8 +776,6 @@ disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[ disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[False-True-Qwen3-8B-FP8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-False-Qwen3-8B-FP8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] -disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] -disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[True-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] disaggregated/test_workers.py::test_workers_conditional_disaggregation[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_workers.py::test_workers_conditional_disaggregation_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] @@ -796,7 +795,6 @@ test_e2e.py::test_eagle3_output_repetition_4gpus[Qwen3/saved_models_Qwen3-235B-A test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] test_e2e.py::test_openai_chat_harmony_perf_metrics test_e2e.py::test_openai_kv_cache_contamination -test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-BF16-llama-3.1-model/Meta-Llama-3.1-8B] test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B-Qwen3/Qwen3-30B-A3B] test_e2e.py::test_ptp_quickstart_advanced_deepseek_r1_8gpus[DeepSeek-R1-DeepSeek-R1/DeepSeek-R1] test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index c460bdda46dc..836c5bd48a0e 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -5,9 +5,6 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_adp16_cuteds test_e2e.py::test_multi_nodes_eval[DeepSeek-R1/DeepSeek-R1-0528-FP4-tp16-mmlu] test_e2e.py::test_multi_nodes_eval[Kimi-K2-Thinking-NVFP4-tp16-mmlu] test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] -test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] -test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] -test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[baseline] accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[dspark] accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[reuse] diff --git a/tests/integration/test_lists/qa/llm_spark_func.yml b/tests/integration/test_lists/qa/llm_spark_func.yml index 446278687b46..8eb91541d395 100644 --- a/tests/integration/test_lists/qa/llm_spark_func.yml +++ b/tests/integration/test_lists/qa/llm_spark_func.yml @@ -11,9 +11,6 @@ llm_spark_func: tests: - test_e2e.py::test_ptp_quickstart_advanced[GPT-OSS-20B-gpt_oss/gpt-oss-20b] - test_e2e.py::test_ptp_quickstart_advanced[GPT-OSS-120B-gpt_oss/gpt-oss-120b] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-bf16-instruct-llama-3.1-model/Llama-3.1-8B-Instruct] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP4-modelopt-hf-model-hub/Llama-3.1-8B-Instruct-fp4] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-8b-fp8-Qwen3/nvidia-Qwen3-8B-FP8] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-8b-nvfp4-Qwen3/nvidia-Qwen3-8B-NVFP4] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-8B-bf16-Qwen3/Qwen3-8B] @@ -34,7 +31,6 @@ llm_spark_func: - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=True] - test_e2e.py::test_trtllm_benchmark_serving[gpt_oss/gpt-oss-20b] - - test_e2e.py::test_openai_chat_guided_decoding[meta-llama/Llama-3.1-8B-Instruct] - examples/serve/test_serve.py::test_nemotron3_super_120b_nvfp4 - examples/serve/test_serve.py::test_nemotron3_nano_omni_nvfp4[text_reasoning_on] - examples/serve/test_serve.py::test_nemotron3_nano_omni_nvfp4[text_reasoning_off] diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5a883a81c248..51e97918e640 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -121,7 +121,6 @@ l0_a10: - test_e2e.py::test_openai_responses_entrypoint - test_e2e.py::test_openai_completions_example[pytorch] - test_e2e.py::test_openai_chat_example[pytorch] TIMEOUT (90) - - test_e2e.py::test_trtllm_bench_request_rate_and_concurrency[enable_concurrency-] - test_e2e.py::test_trtllm_bench_invalid_token_pytorch[TinyLlama-1.1B-Chat-v1.0-TinyLlama-1.1B-Chat-v1.0] # visual_gen - unittest/_torch/visual_gen/test_profiler.py diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d680cce29d4b..b3fa9d799f5e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -81,15 +81,9 @@ l0_b200: - accuracy/test_epd_disagg_multimodal.py::TestVideoMMEEPD::test_disaggregated_videomme[nemotron_nano_v3_omni_nvfp4] - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16_mtp - disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyLlama-1.1B-Chat-v1.0] # nvbugs 5300551 - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - test_e2e.py::test_ptp_quickstart_advanced[Nemotron-Nano-9B-v2-nvfp4-NVIDIA-Nemotron-Nano-9B-v2-NVFP4] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] - test_e2e.py::test_ptp_quickstart_advanced_mtp[DeepSeek-V3-Lite-BF16-DeepSeek-V3-Lite/bf16] - test_e2e.py::test_ptp_quickstart_advanced_mtp_eagle[DeepSeek-V3-Lite-BF16-DeepSeek-V3-Lite/bf16] - - test_e2e.py::test_ptp_quickstart_advanced_mixed_precision - - test_e2e.py::test_ptp_quickstart_advanced_eagle3[Llama-3.1-8b-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct-EAGLE3-LLaMA3.1-Instruct-8B] - - test_e2e.py::test_ptp_quickstart_advanced_ngram[Llama-3.1-8B-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct] - - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False] - test_e2e.py::test_openai_chat_guided_decoding[openai/gpt-oss-120b] # Covers tests/unittest/_torch/attention/. Two sub-trees moved in here from elsewhere # under tests/unittest/_torch/ and have no entry of their own on any list, so this entry diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index 370536bbc609..835f6c8d09ce 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -112,7 +112,6 @@ l0_dgx_h200: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=False] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=True-overlap_scheduler=False-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] - - test_e2e.py::test_trtllm_bench_mgmn - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=False] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.5-fp8kv=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B_Instruct_2507::test_skip_softmax_attention_4gpus[target_sparsity_0.9-fp8kv=False] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 09eef39b35b4..619c89deddf3 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -159,8 +159,6 @@ l0_h100: - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 - accuracy/test_llm_api_pytorch.py::TestLagunaXS_2_1::test_fp8_dflash - accuracy/test_llm_api_pytorch.py::TestNemotron35Lightning::test_nvfp4_marlin_mtp3_chunked_prefill - - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False] - - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-instruct-hf-fp8-True-True] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1_single_gpu_mtp[DeepSeek-V3-Lite-fp8] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_tp1_two_mtp[DeepSeek-V3-Lite-fp8] @@ -175,8 +173,6 @@ l0_h100: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-False-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] - - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[True-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] - - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-False-EAGLE3-LLaMA3.1-Instruct-8B-Llama-3.1-8B-Instruct] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_cancel_gen_requests[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[True-TinyLlama-1.1B-Chat-v1.0] @@ -192,15 +188,10 @@ l0_h100: - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_transcribe_end_to_end - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[bf16-kv-v2-decoder-graphs-on-greedy] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_beam_search[bf16-kv-v1-decoder-graphs-on-beam2] - - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - - test_e2e.py::test_trtllm_bench_iteration_log[PyTorch-non-streaming-meta-llama/Llama-3.1-8B-llama-3.1-model/Meta-Llama-3.1-8B] - - test_e2e.py::test_trtllm_bench_request_rate_and_concurrency[enable_concurrency-enable_request_rate] # negative test - - test_e2e.py::test_trtllm_bench_help_sanity[meta-llama/Llama-3.1-8B] - test_e2e.py::test_openai_chat_harmony - test_e2e.py::test_openai_chat_harmony_perf_metrics - test_e2e.py::test_openai_responses - test_e2e.py::test_anthropic_messages - - test_e2e.py::test_openai_chat_guided_decoding[meta-llama/Llama-3.1-8B-Instruct] - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph[v1] - kv_cache/test_final_single_token_context_cuda_graph.py::test_final_token_reuse_cuda_graph[v2] - kv_cache/test_final_single_token_context_cuda_graph.py::test_changed_final_token_reuse_cuda_graph[v1] @@ -236,7 +227,6 @@ l0_h100: tests: - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales[mtp=disable-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True] - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_fp8_block_scales[latency] - - test_e2e.py::test_trtllm_bench_pytorch_backend_sanity[meta-llama/Llama-3.1-8B-llama-3.1-8b-False-False] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_llama[False-False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_deepseek[False-False-DeepSeek-V3-Lite-fp8/fp8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[False-False-Qwen3-8B-FP8] @@ -408,8 +398,6 @@ l0_h100: - accuracy/test_llm_api_pytorch_multimodal.py::TestMistralSmall24B::test_auto_dtype[forced_chunked_prefill] - accuracy/test_llm_api_pytorch_multimodal.py::TestQwen3VL::test_auto_dtype[forced_chunked_prefill] - accuracy/test_llm_api_pytorch_multimodal.py::TestNemotron_Nano_12B_V2_VL::test_auto_dtype[forced_chunked_prefill] - - test_e2e.py::test_draft_token_tree_quickstart_advanced_eagle3[Llama-3.1-8b-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct-EAGLE3-LLaMA3.1-Instruct-8B] - - test_e2e.py::test_draft_token_tree_quickstart_advanced_eagle3_depth_1_tree[Llama-3.1-8b-Instruct-llama-3.1-model/Llama-3.1-8B-Instruct-EAGLE3-LLaMA3.1-Instruct-8B] # ---- moved to post-merge (MoE CI optimization) ---- - unittest/_torch/moe/test_moe_backend.py::test_moe_backend -k "CUTLASS" # ---- non-quantized (quant=None) moved to post-merge ---- diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index 7e94f1147991..6cd4b222b5cb 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -29,9 +29,6 @@ l0_rtx_pro_6000: - unittest/_torch/thop/parallel/test_w4a8_linear.py - test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-BF16-llama-3.1-model/Meta-Llama-3.1-8B] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-NVFP4-nvfp4-quantized/Meta-Llama-3.1-8B] - - test_e2e.py::test_ptp_quickstart_advanced[Llama3.1-8B-FP8-llama-3.1-model/Llama-3.1-8B-Instruct-FP8] - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B-Qwen3/Qwen3-30B-A3B] # 3mins - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B_fp8_hf-Qwen3/saved_models_Qwen3-30B-A3B_fp8_hf] # 3mins - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B_nvfp4_hf-Qwen3/saved_models_Qwen3-30B-A3B_nvfp4_hf] # 2mins diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index cc7a3cdd7ec1..56f8d1e5097e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -92,7 +92,6 @@ examples/visual_gen/test_visual_gen_wan.py::test_wan_feature_accuracy_against_go full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:A100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:A100/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) -full:A100/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:A100/disaggregated/test_disaggregated.py::test_disaggregated_python_transceiver_host_offload[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6758573) full:A100/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6758594) full:A100/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6758594) @@ -114,7 +113,6 @@ full:B200/accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpu full:B200/accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm] SKIP (https://nvbugs/6731978) full:B200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:B200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) -full:B200/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-gpt_oss_120b_eagle_trtllm_stress] SKIP (https://nvbugs/6649384) full:B200/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6621362) @@ -130,7 +128,6 @@ full:B300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4 full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp4_gentp4[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) -full:B300/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:DGX_B200/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16[mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-enable_chunked_prefill=False-v2_kv_cache=True] SKIP (https://nvbugs/6633268) full:DGX_B200/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py::test_kda_mixer_empty_prefill SKIP (https://nvbugs/6705034) @@ -155,20 +152,22 @@ full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_ full:GB200/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:GB200/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) -full:GB200/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:GB300/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_first[adp-mtp2] SKIP (https://nvbugs/6295740) full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=2-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=True] SKIP (https://nvbugs/6432818) full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=False-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False-sampler_async_worker=False] SKIP (https://nvbugs/6661948) full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6661948) full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6697099) full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] SKIP (https://nvbugs/6714109) +full:GB300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6572838) +full:GB300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm_eagle] SKIP (https://nvbugs/6572838) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (https://nvbugs/6487918) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_fp8_tp16_mtp3_trtllm SKIP (https://nvbugs/6694922) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) -full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) +full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[block_reuse-python_scheduler] SKIP (https://nvbugs/6608387) +full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[no_reuse-python_scheduler] SKIP (https://nvbugs/6608387) full:GB300/unittest/_torch/modeling/test_modeling_gpt_oss.py::test_gpt_oss_trtllmgen[CUTLASS] SKIP (https://nvbugs/6633932) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6313072) @@ -176,7 +175,6 @@ full:H100/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_dummy_load_format S full:H100/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp1-CUTLASS] SKIP (https://nvbugs/6732067) full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:H100/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) -full:H100/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:H100/disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-qwen3_32b_fp8_stress] SKIP (https://nvbugs/6312828) full:H100/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_type_default SKIP (https://nvbugs/6727262) full:H100/test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B-Qwen3/Qwen3-30B-A3B] SKIP (https://nvbugs/6700875) @@ -199,13 +197,11 @@ full:H20/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[t full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[bf16] SKIP (https://nvbugs/6618649) -full:H20/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:H20/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6692009) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) -full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_logging SKIP (https://nvbugs/6727262) full:L40S/llmapi/test_llm_api_qa.py::TestLlmDefaultBackend::test_llm_args_type_default SKIP (https://nvbugs/6727262) @@ -266,6 +262,9 @@ test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/ test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (bug pending, tracked in PR 17414) unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) unittest/_torch/attention/sparse/msa/test_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) +unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype0] SKIP (https://nvbugs/6669902) +unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype1] SKIP (https://nvbugs/6669902) +unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py::test_sparse_decode_matches_msa_kernel SKIP (https://nvbugs/6669902) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[exaone_moe_gqa_swa128-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6668773) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6641268) unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy SKIP (https://nvbugs/6607482) @@ -277,7 +276,6 @@ unittest/_torch/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048 unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:16-seqlen:2] SKIP (https://nvbugs/6501404) unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[False] SKIP (https://nvbugs/6535767) unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[True] SKIP (https://nvbugs/6535767) -unittest/_torch/speculative/hw_agnostic/test_ngram.py::test_llama_ngram[True-True-TRTLLM] SKIP (https://nvbugs/6507102) unittest/_torch/speculative/test_eagle3.py SKIP (https://nvbugs/5461761) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cublaslt] SKIP (https://nvbugs/6581067) unittest/_torch/thop/parallel/test_fp4_linear.py::test_fp4_gemm_bias_per_backend[mnk2-cutlass] SKIP (https://nvbugs/6581067) diff --git a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py index 4b93cd54dcea..106a30b0ac14 100644 --- a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py @@ -1,14 +1,11 @@ -import json import math -import os import pytest import torch -from utils.llm_data import llm_models_root from utils.util import getSMVersion import tensorrt_llm -from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm import SamplingParams from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.sparse.rocket import ( RocketKVCacheManager, @@ -21,86 +18,10 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes from tensorrt_llm.bindings import DataType -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, RocketSparseAttentionConfig +from tensorrt_llm.llmapi import KvCacheConfig, RocketSparseAttentionConfig from tensorrt_llm.mapping import Mapping -@pytest.mark.skipif(getSMVersion() < 100, reason="RocketKV requires SM100 (Blackwell)") -@pytest.mark.parametrize("backend", ["pytorch"]) -@pytest.mark.parametrize("model_name", ["llama-3.1-model/Llama-3.1-8B-Instruct"]) -@pytest.mark.parametrize("attention_backend", ["VANILLA", "TRTLLM"]) -def test_model(backend, model_name, attention_backend, monkeypatch): - # RocketKV is a single-GPU path. Keep this test independent of MPI - # dynamic-process bootstrap so a cluster launch failure cannot mask the - # attention result. - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - model_dir = str(llm_models_root() / model_name) - max_batch_size = 16 - max_output_tokens = 128 - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, enable_block_reuse=False) - - kt_cache_dtype = "float8_e5m2" if attention_backend == "TRTLLM" else "bfloat16" - - sparse_attention_config = RocketSparseAttentionConfig( - window_size=32, - kernel_size=63, - prompt_budget=2048, - kt_cache_dtype=kt_cache_dtype, - ) - - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1, 2, 4, 8, 16], - enable_padding=True, - ) - - llm = LLM( - model=model_dir, - backend=backend, - kv_cache_config=kv_cache_config, - attn_backend=attention_backend, - sparse_attention_config=sparse_attention_config, - max_batch_size=max_batch_size, - max_seq_len=20480, - max_num_tokens=81920, - cuda_graph_config=None if attention_backend == "VANILLA" else cuda_graph_config, - ) - - inputs, references = [], [] - current_file = os.path.abspath(__file__) - current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) - input_file = f"{current_dir}/multi_gpu/NIAH_simple_data.jsonl" - with open(input_file, "r") as f: - for line in f: - sample = json.loads(line) - inputs.append( - { - "prompt": sample["input_context"] + sample["input_query"], - } - ) - references.append(sample["outputs"][0]) - - with llm: - outputs = llm.generate( - inputs, - use_tqdm=True, - sampling_params=SamplingParams( - add_special_tokens=False, max_tokens=max_output_tokens, temperature=0.8, top_p=0.95 - ), - ) - - count = 0 - for ref, ret in zip(references, outputs): - print(f"ret: {ret.outputs[0].text}") - print(f"ref: {ref}") - if ref not in ret.outputs[0].text: - print(f"reference {ref} is not in the output {ret.outputs[0].text}") - else: - count = count + 1 - acc = count / len(outputs) - - assert acc >= 0.9, "accuracy test of rocketkv sparse attention failed" - - def create_rocket_kv_cache_manager( num_layers, num_kv_heads, @@ -772,11 +693,6 @@ def test_rocket_add_dummy_requests_forwards_capture_sampling_params(mocker): if __name__ == "__main__": - # RocketKV e2e tests - print("=== Testing RocketKV E2E tests ===") - test_model("pytorch", "llama-3.1-model/Llama-3.1-8B-Instruct", "VANILLA") - test_model("pytorch", "llama-3.1-model/Llama-3.1-8B-Instruct", "TRTLLM") - # Unit tests for sparse_kv_predict print("\n=== Testing sparse_kv_predict ===") test_sparse_kv_predict(1, 1) # bs=1, context only diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index c3fa5f427a5f..c02e04053f20 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -1,66 +1,7 @@ -import unittest -from copy import deepcopy from dataclasses import dataclass from typing import Any import torch -from _torch.helpers import create_mock_cuda_graph_runner -from parameterized import parameterized -from transformers import LlamaConfig -from transformers import LlamaForCausalLM as HFLlamaForCausalLM -from utils.llm_data import llm_models_root -from utils.util import default_dtype, getSMVersion - -import tensorrt_llm -from tensorrt_llm._torch.attention.backends.utils import get_attention_backend -from tensorrt_llm._torch.metadata import KVCacheParams -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_llama import LlamaForCausalLM -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ - _update_kv_cache_draft_token_location -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager -from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._torch.speculative.interface import (SpecMetadata, - SpeculativeDecodingMode) -from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager -from tensorrt_llm._utils import get_sm_version -from tensorrt_llm.bindings.executor import KvCacheConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.modeling_utils import QuantConfig - -LLAMA_3_1_8B_CONFIG = { - "architectures": ["LlamaForCausalLM"], - "attention_bias": False, - "attention_dropout": 0.0, - "bos_token_id": 128000, - "eos_token_id": 128001, - "hidden_act": "silu", - "hidden_size": 4096, - "initializer_range": 0.02, - "intermediate_size": 14336, - "max_position_embeddings": 131072, - "mlp_bias": False, - "model_type": "llama", - "num_attention_heads": 32, - "num_hidden_layers": 32, - "num_key_value_heads": 8, - "pretraining_tp": 1, - "rms_norm_eps": 1e-05, - "rope_scaling": { - "factor": 8.0, - "low_freq_factor": 1.0, - "high_freq_factor": 4.0, - "original_max_position_embeddings": 8192, - "rope_type": "llama3" - }, - "rope_theta": 500000.0, - "tie_word_embeddings": False, - "torch_dtype": "bfloat16", - "transformers_version": "4.43.0.dev0", - "use_cache": True, - "vocab_size": 128256 -} @dataclass(repr=False) @@ -82,658 +23,3 @@ def reduce_llama_config(mem_for_full_model: int, num_layers = int(config_dict["num_hidden_layers"] * model_fraction) num_layers = min(num_layers, default_num_layers) config_dict["num_hidden_layers"] = num_layers - - -class TestLlama(unittest.TestCase): - - @parameterized.expand([None, "FP8"]) - def test_llama_sanity(self, quant_algo): - config_dict = deepcopy(LLAMA_3_1_8B_CONFIG) - # 8B * sizeof(float16) plus some extra for activations - mem_for_full_model = (2 + 1) * 8 * 2**(30) - reduce_llama_config(mem_for_full_model, config_dict) - if config_dict["num_hidden_layers"] <= 0: - self.skipTest("Insufficient memory for a single Llama layer") - llama_config = LlamaConfig.from_dict(config_dict) - if quant_algo: - quant_config = QuantConfig(quant_algo=quant_algo) - else: - quant_config = None - if quant_algo == "FP8" and getSMVersion() < 89: - self.skipTest("This test is not supported in pre-Ada architecture") - - dtype = llama_config.torch_dtype - device = torch.device('cuda') - - with torch.device(device), default_dtype(dtype): - model_config = ModelConfig(pretrained_config=llama_config, - quant_config=quant_config) - llama = LlamaForCausalLM(model_config).to(device) - - input_ids = torch.tensor([100, 200, 300, 100, 200, 100, 400, 500], - dtype=torch.int, - device=device) - - context_sequence_lengths = [3, 2, 1] - sequence_lengths = context_sequence_lengths + [1, 1] - past_seen_tokens = [0, 0, 0, 62, 75] - request_ids = list(range(len(sequence_lengths))) - token_nums = (torch.tensor(past_seen_tokens) + - torch.tensor(sequence_lengths)).tolist() - prompt_lens = token_nums[:3] + past_seen_tokens[3:] - - num_blocks = 100 - tokens_per_block = 128 - head_dim = llama.config.hidden_size // llama.config.num_attention_heads - num_layers = llama.config.num_hidden_layers - num_kv_heads = llama.config.num_key_value_heads - max_seq_len = num_blocks * tokens_per_block - batch_size = len(context_sequence_lengths) + 2 - - if dtype == torch.half: - kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF - elif dtype == torch.bfloat16: - kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 - else: - raise ValueError("Invalid dtype") - - mapping = Mapping(world_size=1, tp_size=1, rank=0) - kv_cache_config = KvCacheConfig(max_tokens=num_blocks * - tokens_per_block) - kv_cache_manager = KVCacheManager( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=batch_size, - mapping=mapping, - dtype=kv_cache_dtype, - ) - kv_cache_manager.add_dummy_requests(request_ids, token_nums) - - metadata_cls = get_attention_backend(model_config.attn_backend).Metadata - attn_metadata = metadata_cls( - seq_lens=torch.tensor(sequence_lengths, dtype=torch.int), - num_contexts=len(context_sequence_lengths), - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=past_seen_tokens, - ), - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - max_num_requests=len(context_sequence_lengths) + 2, - max_num_tokens=8192, - ) - - position_ids = [] - for i, tokens in enumerate(past_seen_tokens): - seq_len = context_sequence_lengths[i] if i < len( - context_sequence_lengths) else 1 - position_id = torch.arange(tokens, - tokens + seq_len, - device=input_ids.device) - position_ids.append(position_id) - - position_ids = torch.cat(position_ids).unsqueeze(0) - - with torch.inference_mode(): - attn_metadata.prepare() - logits = llama.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata) - - self.assertEqual(len(past_seen_tokens), logits.shape[0]) - - with torch.inference_mode(): - attn_metadata.prepare() - logits = llama.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - return_context_logits=True) - self.assertEqual(input_ids.shape, logits.shape[:-1]) - - kv_cache_manager.shutdown() - - @parameterized.expand([ - Scenario(backend="VANILLA"), - Scenario(backend="FLASHINFER"), - Scenario(backend="FLASHINFER", use_cuda_graph=True), - Scenario(backend="TRTLLM"), - Scenario(backend="TRTLLM", use_cuda_graph=True), - ], lambda testcase_func, param_num, param: - f"{testcase_func.__name__}[{param.args[0]}]") - @torch.no_grad() - def test_llama_allclose_to_hf(self, scenario: Scenario) -> None: - """ - Compare output to HF - """ - backend = scenario.backend - metadata_cls = get_attention_backend(backend).Metadata - - torch.random.manual_seed(0) - config_dict = deepcopy(LLAMA_3_1_8B_CONFIG) - # 8B * sizeof(float16) plus some extra for activations - # times 2, since we'll need 2 of these - mem_for_full_model = (2 + 1) * 8 * 2**(30) * 4 - reduce_llama_config(mem_for_full_model, config_dict) - if config_dict["num_hidden_layers"] <= 0: - self.skipTest("Insufficient memory for a single Llama layer") - llama_config = LlamaConfig.from_dict(config_dict) - dtype = llama_config.torch_dtype - device = torch.device('cuda') - - with torch.device(device), default_dtype(dtype): - hf_llama = HFLlamaForCausalLM(llama_config).eval() - - model_config = ModelConfig(pretrained_config=llama_config, - attn_backend=backend) - - llama = LlamaForCausalLM(model_config).to(dtype).to(device) - llama.load_weights(hf_llama.state_dict()) - llama.post_load_weights() - - num_blocks = 1 - tokens_per_block = 128 - head_dim = llama.config.hidden_size // llama.config.num_attention_heads - num_layers = llama.config.num_hidden_layers - num_kv_heads = llama.config.num_key_value_heads - max_seq_len = num_blocks * tokens_per_block - batch_size = 1 - - if dtype == torch.half: - kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF - elif dtype == torch.bfloat16: - kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 - else: - raise ValueError("Invalid dtype") - - mapping = Mapping(world_size=1, tp_size=1, rank=0) - kv_cache_config = KvCacheConfig(max_tokens=num_blocks * - tokens_per_block) - kv_cache_manager = KVCacheManager( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=batch_size, - mapping=mapping, - dtype=kv_cache_dtype, - ) - - # context - input_ids = torch.tensor([100, 200, 300, 100, 200, 100, 400, 500], - dtype=torch.int, - device=device) - - num_cached_tokens_per_seq = [0] - request_ids = [1] - token_nums = [input_ids.size(-1)] - prompt_lens = [input_ids.size(-1)] - kv_cache_manager.add_dummy_requests(request_ids, token_nums) - - attn_metadata = metadata_cls( - seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), - num_contexts=1, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - ) - - # Note: no CUDA graphs for prefill, the graph runner is built for - # decoding only. - position_ids = [torch.arange(0, input_ids.size(-1))] - position_ids = torch.cat(position_ids).unsqueeze(0).cuda() - with torch.inference_mode(): - attn_metadata.prepare() - logits = llama.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata) - ref = hf_llama.forward(input_ids=input_ids.unsqueeze(0), - position_ids=position_ids, - use_cache=True) - - torch.testing.assert_close(logits, - ref.logits[:, -1].float(), - atol=0.4, - rtol=0.4) - - # gen - gen_input_ids = torch.tensor([600], dtype=torch.int, device=device) - - num_cached_tokens_per_seq = [input_ids.size(-1)] - - attn_metadata = metadata_cls( - seq_lens=torch.tensor([gen_input_ids.size(-1)], dtype=torch.int), - num_contexts=0, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - ) - - gen_position_ids = [ - torch.arange(input_ids.size(-1), - input_ids.size(-1) + gen_input_ids.size(-1)) - ] - gen_position_ids = torch.cat(gen_position_ids).unsqueeze(0).cuda() - - graph_runner = create_mock_cuda_graph_runner( - 1) if scenario.use_cuda_graph else None - - def run_forward(input_ids, position_ids, attn_metadata): - attn_metadata.prepare() - if not scenario.use_cuda_graph: - return llama.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata) - else: - inputs = { - "input_ids": input_ids, - "position_ids": position_ids, - "attn_metadata": attn_metadata, - } - key = (1, 0, False) - graph_runner.capture(key, - lambda inputs: llama.forward(**inputs), - inputs) - for _ in range(2): - # Run it twice. This helps us catch problems if buffers are accidentally reallocated - # in prepare(). - attn_metadata.prepare() - logits = graph_runner.replay(key, inputs) - return logits - - if scenario.use_cuda_graph: - attn_metadata = attn_metadata.create_cuda_graph_metadata(1) - - with torch.inference_mode(): - logits = run_forward(input_ids=gen_input_ids, - position_ids=gen_position_ids, - attn_metadata=attn_metadata) - ref = hf_llama.forward(input_ids=gen_input_ids.unsqueeze(0), - position_ids=gen_position_ids, - past_key_values=ref.past_key_values, - use_cache=True) - - torch.testing.assert_close(logits, - ref.logits[:, -1].float(), - atol=0.4, - rtol=0.4) - if graph_runner is not None: - graph_runner.clear() - kv_cache_manager.shutdown() - - @torch.no_grad() - def test_llama_verification_with_kv_cache_relocation(self) -> None: - """ - Verify the output of the model with kv cache relocation - """ - backend = "TRTLLM" - metadata_cls = get_attention_backend(backend).Metadata - - config_dict = deepcopy(LLAMA_3_1_8B_CONFIG) - - llama_config = LlamaConfig.from_dict(config_dict) - dtype = llama_config.torch_dtype - device = torch.device('cuda') - - with torch.device(device), default_dtype(dtype): - models_path = llm_models_root() - model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - hf_llama = HFLlamaForCausalLM.from_pretrained( - model_dir, - torch_dtype=torch.bfloat16, - device_map="cuda", - ).eval() - - model_config = ModelConfig(pretrained_config=llama_config, - attn_backend=backend) - - llama = LlamaForCausalLM(model_config).to(dtype).to(device) - llama.load_weights(hf_llama.state_dict()) - - num_blocks = 2 - tokens_per_block = 32 - head_dim = llama.config.hidden_size // llama.config.num_attention_heads - num_layers = llama.config.num_hidden_layers - num_kv_heads = llama.config.num_key_value_heads - num_heads_per_kv = llama.config.num_attention_heads // num_kv_heads - max_seq_len = num_blocks * tokens_per_block - batch_size = 1 - - if dtype == torch.half: - kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF - elif dtype == torch.bfloat16: - kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 - else: - raise ValueError("Invalid dtype") - kv_cache_dtype_byte_size = 2 - - mapping = Mapping(world_size=1, tp_size=1, rank=0) - kv_cache_config = KvCacheConfig(max_tokens=num_blocks * - tokens_per_block) - kv_cache_manager = KVCacheManager( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=batch_size, - mapping=mapping, - dtype=kv_cache_dtype, - ) - - # context - input_ids = torch.tensor([ - 128000, 32, 6369, 1990, 264, 22999, 1217, 323, 459, 21075, 11478, - 18328, 13, 578, 18328, 6835, 11190, 11, 11944, 11, 323, 48887, - 11503, 311, 279, 1217, 596, 4860, 13, 14194, 25, 22691, 36660, 3931, - 2891, 25 - ], - dtype=torch.int, - device=device) - - num_cached_tokens_per_seq = [0] - request_ids = [900] - token_nums = [input_ids.size(-1)] - prompt_lens = [input_ids.size(-1)] - requests = kv_cache_manager.add_dummy_requests(request_ids, token_nums) - request = requests[0] - - attn_metadata = metadata_cls( - seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), - num_contexts=1, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - num_heads_per_kv=num_heads_per_kv, - ) - - position_ids = [torch.arange(0, input_ids.size(-1))] - position_ids = torch.cat(position_ids).unsqueeze(0).cuda() - with torch.inference_mode(): - attn_metadata.prepare() - logits = llama.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata) - - def run_forward(input_ids, position_ids, attn_metadata): - return llama.forward(input_ids=input_ids, - position_ids=position_ids, - attn_metadata=attn_metadata, - return_context_logits=True) - - # prepare for the first generation - gen_input_ids_0 = torch.tensor([ - 22691, 11, 0, 13, 15592, 323, 315, 12, 311, 362, 220, 32, 362, 426, - 330, 358, 362, 358, 358, 362, 32, 0, 13, 32, 6369 - ], - dtype=torch.int, - device=device) - spec_decoding_position_offsets = torch.tensor([ - 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, - 2, 3 - ], - dtype=torch.int, - device=device) - spec_decoding_packed_mask = torch.tensor( - [ - 1, 3, 5, 9, 17, 33, 65, 129, 257, 513, 1025, 2051, 4099, 8195, - 16387, 32771, 65541, 131077, 262153, 524297, 1048593, 2097169, - 4194321, 8388641, 16842757 - ], - dtype=torch.int, - device=device).unsqueeze(0).unsqueeze(2) - - num_cached_tokens_per_seq = [input_ids.size(-1)] - is_spec_decoding_enabled = True - use_spec_decoding = True - is_spec_dec_tree = True - is_spec_dec_dynamic_tree = True - max_total_draft_tokens = gen_input_ids_0.size(-1) - 1 - - spec_tree_mgr = SpecTreeManager( - max_num_requests=1, - use_dynamic_tree=is_spec_dec_dynamic_tree, - max_total_draft_tokens=max_total_draft_tokens, - max_draft_len=max_total_draft_tokens, - eagle_choices=None, - dynamic_tree_max_topK=10, - ) - # Populate with test data - spec_tree_mgr.spec_dec_position_offsets[:1, :].copy_( - spec_decoding_position_offsets.unsqueeze(0)) - spec_tree_mgr.spec_dec_packed_mask[:1, :, :].copy_( - spec_decoding_packed_mask) - - attn_metadata_gen_phase_0 = metadata_cls( - seq_lens=torch.tensor([gen_input_ids_0.size(-1)], dtype=torch.int), - num_contexts=0, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - is_spec_decoding_enabled=is_spec_decoding_enabled, - use_spec_decoding=use_spec_decoding, - is_spec_dec_tree=is_spec_dec_tree, - is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, - num_heads_per_kv=num_heads_per_kv, - ) - attn_metadata_gen_phase_0.prepare() - attn_metadata_gen_phase_0.update_spec_dec_param( - batch_size=batch_size, - is_spec_decoding_enabled=is_spec_decoding_enabled, - is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, - is_spec_dec_tree=is_spec_dec_tree, - max_draft_len=max_total_draft_tokens, - max_total_draft_tokens=max_total_draft_tokens, - spec_tree_manager=spec_tree_mgr, - ) - - gen_position_ids_0 = [ - torch.full((gen_input_ids_0.size(-1), ), - input_ids.size(-1), - dtype=torch.int64) - ] - gen_position_ids_0 = torch.cat(gen_position_ids_0).unsqueeze(0).cuda() - - with torch.inference_mode(): - gen_logits_0 = run_forward(input_ids=gen_input_ids_0, - position_ids=gen_position_ids_0, - attn_metadata=attn_metadata_gen_phase_0) - - request.py_num_accepted_draft_tokens = 1 - request.py_num_accepted_draft_tokens_indices = [1] - request.py_rewind_len = gen_input_ids_0.size( - -1) - request.py_num_accepted_draft_tokens - 1 - request.state = LlmRequestState.GENERATION_IN_PROGRESS - scheduled_requests = ScheduledRequests() - scheduled_requests.generation_requests = [request] - kv_cache_manager.max_draft_len = gen_input_ids_0.size(-1) - 1 - _update_kv_cache_draft_token_location(kv_cache_manager, - scheduled_requests, - attn_metadata_gen_phase_0, - kv_cache_dtype_byte_size) - if request.py_rewind_len > 0: - kv_cache_manager.rewind_kv_cache(request, request.py_rewind_len) - torch.cuda.synchronize() - - # prepare for the second generation - gen_input_ids_1 = torch.tensor([2650, 649], - dtype=torch.int, - device=device) - - num_cached_tokens_per_seq_1 = [ - input_ids.size(-1) + request.py_num_accepted_draft_tokens + 1 - ] - attn_metadata_gen_phase_0.seq_lens = torch.tensor( - [gen_input_ids_1.size(-1)], dtype=torch.int) - attn_metadata_gen_phase_0.kv_cache_params.num_cached_tokens_per_seq = num_cached_tokens_per_seq_1 - - attn_metadata_gen_phase_0.spec_decoding_position_offsets = None - attn_metadata_gen_phase_0.spec_decoding_packed_mask = None - attn_metadata_gen_phase_0.spec_decoding_generation_lengths = None - attn_metadata_gen_phase_0.prepare() - is_tree_phase1 = is_spec_dec_tree if get_sm_version() < 100 else False - spec_tree_mgr_phase1 = None - spec_metadata_phase1 = None - if is_tree_phase1: - max_draft_1 = gen_input_ids_1.size(-1) - 1 - spec_tree_mgr_phase1 = SpecTreeManager( - max_num_requests=1, - use_dynamic_tree=True, - max_total_draft_tokens=max_draft_1, - max_draft_len=max_draft_1, - eagle_choices=None, - dynamic_tree_max_topK=10, - ) - spec_metadata_phase1 = SpecMetadata( - max_num_requests=1, - max_draft_len=max_draft_1, - max_total_draft_tokens=max_draft_1, - spec_dec_mode=SpeculativeDecodingMode.EAGLE3, - ) - attn_metadata_gen_phase_0.update_spec_dec_param( - batch_size=batch_size, - is_spec_decoding_enabled=is_spec_decoding_enabled, - is_spec_dec_tree=is_tree_phase1, - is_spec_dec_dynamic_tree=is_tree_phase1, - max_draft_len=gen_input_ids_1.size(-1) - 1, - max_total_draft_tokens=gen_input_ids_1.size(-1) - 1, - spec_metadata=spec_metadata_phase1, - spec_tree_manager=spec_tree_mgr_phase1) - - gen_position_ids_1 = [ - torch.full( - (gen_input_ids_1.size(-1), ), - input_ids.size(-1) + request.py_num_accepted_draft_tokens + 1, - dtype=torch.int64) - ] - gen_position_ids_1 = torch.cat(gen_position_ids_1).unsqueeze(0).cuda() - - with torch.inference_mode(): - gen_logits_1 = run_forward(input_ids=gen_input_ids_1, - position_ids=gen_position_ids_1, - attn_metadata=attn_metadata_gen_phase_0) - - torch.cuda.synchronize() - - # prepare for the reference generation - gen_input_ids_ref = torch.tensor([22691, 0, 2650, 649], - dtype=torch.int, - device=device) - num_cached_tokens_per_seq_ref = [input_ids.size(-1)] - - attn_metadata_ref = metadata_cls( - seq_lens=torch.tensor([gen_input_ids_ref.size(-1)], - dtype=torch.int), - num_contexts=0, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=num_cached_tokens_per_seq_ref, - ), - max_num_requests=1, - max_num_tokens=8192, - kv_cache_manager=kv_cache_manager, - request_ids=request_ids, - prompt_lens=prompt_lens, - is_spec_decoding_enabled=is_spec_decoding_enabled, - use_spec_decoding=use_spec_decoding, - is_spec_dec_tree=is_spec_dec_tree, - is_spec_dec_dynamic_tree=False, - num_heads_per_kv=num_heads_per_kv, - ) - - attn_metadata_ref.spec_decoding_position_offsets = None - attn_metadata_ref.spec_decoding_packed_mask = None - attn_metadata_ref.spec_decoding_generation_lengths = None - attn_metadata_ref.prepare() - is_tree_ref = is_spec_dec_tree if get_sm_version() < 100 else False - spec_tree_mgr_ref = None - spec_metadata_ref = None - if is_tree_ref: - max_draft_ref = gen_input_ids_ref.size(-1) - 1 - spec_tree_mgr_ref = SpecTreeManager( - max_num_requests=1, - use_dynamic_tree=True, - max_total_draft_tokens=max_draft_ref, - max_draft_len=max_draft_ref, - eagle_choices=None, - dynamic_tree_max_topK=10, - ) - spec_metadata_ref = SpecMetadata( - max_num_requests=1, - max_draft_len=max_draft_ref, - max_total_draft_tokens=max_draft_ref, - spec_dec_mode=SpeculativeDecodingMode.EAGLE3, - ) - attn_metadata_ref.update_spec_dec_param( - batch_size=batch_size, - is_spec_decoding_enabled=is_spec_decoding_enabled, - is_spec_dec_tree=is_tree_ref, - is_spec_dec_dynamic_tree=is_tree_ref, - max_draft_len=gen_input_ids_ref.size(-1) - 1, - max_total_draft_tokens=gen_input_ids_ref.size(-1) - 1, - spec_metadata=spec_metadata_ref, - spec_tree_manager=spec_tree_mgr_ref) - - gen_position_ids_ref = [ - torch.full((gen_input_ids_ref.size(-1), ), - input_ids.size(-1), - dtype=torch.int64) - ] - gen_position_ids_ref = torch.cat(gen_position_ids_ref).unsqueeze( - 0).cuda() - - with torch.inference_mode(): - gen_logits_ref = run_forward(input_ids=gen_input_ids_ref, - position_ids=gen_position_ids_ref, - attn_metadata=attn_metadata_ref) - - torch.cuda.synchronize() - torch.testing.assert_close(gen_logits_1[0, :], - gen_logits_ref[2, :], - atol=1.0, - rtol=1.0) - torch.testing.assert_close(gen_logits_1[1, :], - gen_logits_ref[3, :], - atol=1.0, - rtol=1.0) - - token_id_ref = torch.argmax(gen_logits_ref[3, :], dim=-1) - token_id_gen = torch.argmax(gen_logits_1[1, :], dim=-1) - assert token_id_ref == token_id_gen, "Greedy sampling token id not match" - - kv_cache_manager.shutdown() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py deleted file mode 100644 index 1ac2e3a5948a..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py +++ /dev/null @@ -1,84 +0,0 @@ -# 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 unittest - -import pytest -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.llmapi import CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig - - -@pytest.mark.high_cuda_memory -def test_eagle3_one_model_kv_cache_reuse() -> None: - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - kv_cache_config = KvCacheConfig( - enable_block_reuse=True, - enable_partial_reuse=True, - use_kv_cache_manager_v2=True, - tokens_per_block=32, - max_tokens=8192, - ) - llm_common_config = dict( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=True, - cuda_graph_config=CudaGraphConfig(batch_sizes=[1]), - max_batch_size=1, - kv_cache_config=kv_cache_config, - max_seq_len=8192, - ) - llm_spec = LLM( - **llm_common_config, - speculative_config=Eagle3DecodingConfig( - max_draft_len=4, - speculative_model=eagle_model_dir, - eagle3_one_model=True, - ), - ) - prompt = ("The quick brown fox jumped over the lazy dog. " * 20).strip() + " Once upon a time," - poison_prompt = "X Y Z W A B C D E F G H I J K L M N O P Q R S T U V " * 60 - sampling_params = SamplingParams( - max_tokens=64, - temperature=0, - ignore_eos=True, - return_perf_metrics=True, - ) - - llm_spec.generate(poison_prompt, SamplingParams(max_tokens=8, temperature=0)) - cold_result = llm_spec.generate(prompt, sampling_params) - reuse_result = llm_spec.generate(prompt, sampling_params) - llm_spec.shutdown() - - cold_acceptance = float(cold_result.avg_decoded_tokens_per_iter) - reuse_acceptance = float(reuse_result.avg_decoded_tokens_per_iter) - assert cold_result.outputs[0].text == reuse_result.outputs[0].text - assert cold_acceptance > 1.0 - assert reuse_acceptance > 1.0 - assert reuse_result.cached_tokens >= 64 - assert reuse_acceptance >= cold_acceptance - 0.2 - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py deleted file mode 100644 index adc87b4f23f8..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py +++ /dev/null @@ -1,68 +0,0 @@ -import unittest - -import pytest -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, NGramDecodingConfig - - -# TODO: add disable_overlap_scheduler=False -@pytest.mark.parametrize( - "disable_overlap_scheduler,use_cuda_graph,attn_backend", - [[True, False, "TRTLLM"], [True, True, "TRTLLM"], [True, False, "FLASHINFER"]], -) -@pytest.mark.high_cuda_memory -def test_llama_ngram(disable_overlap_scheduler: bool, use_cuda_graph: bool, attn_backend: str): - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 20: - pytest.skip("Not enough memory to load target model") - - max_batch_size = 2 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) - cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) if use_cuda_graph else None - - llm_common_config = dict( - model=llm_models_root() / "llama-3.1-model" / "Meta-Llama-3.1-8B", - backend="pytorch", - attn_backend=attn_backend, - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_num_tokens=2048, - ) - - spec_config = NGramDecodingConfig( - max_draft_len=max_draft_len, - max_matching_ngram_size=2, - is_keep_all=True, - is_use_oldest=True, - is_public_pool=True, - ) - - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) - - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - results_spec = llm_spec.generate(prompts, sampling_params) - generated_text_spec = [result.outputs[0].text for result in results_spec] - llm_spec.shutdown() - - llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) - generated_text_ref = [result.outputs[0].text for result in results_ref] - llm_ref.shutdown() - - for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): - # The spec decode algorithm currently guarantees identical results - assert text_spec == text_ref - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py deleted file mode 100644 index 5e93a70b174c..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py +++ /dev/null @@ -1,115 +0,0 @@ -import unittest - -import pytest -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, PARDDecodingConfig - - -@pytest.mark.parametrize("disable_overlap_scheduler", [True, False]) -def test_pard(disable_overlap_scheduler: bool): - """Test PARD speculative decoding with CUDA graph support. - - This test verifies that PARD (Parallel Draft) speculative decoding works - correctly with CUDA graphs and padding enabled. - """ - attn_backend = "TRTLLM" - enable_block_reuse = False - enable_chunked_prefill = False - - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - # Test with 3 requests and max_batch_size=4 to trigger padding - max_batch_size = 4 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, max_tokens=2048) - use_cuda_graph = True - cuda_graph_config = ( - CudaGraphConfig(batch_sizes=[1, 2, 4], enable_padding=True) if use_cuda_graph else None - ) - - llm_common_config = dict( - model=target_model_dir, - attn_backend=attn_backend, - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_seq_len=2048, - enable_chunked_prefill=enable_chunked_prefill, - ) - - spec_config = PARDDecodingConfig( - max_draft_len=max_draft_len, - speculative_model=pard_model_dir, - ) - - # Create the LLM instance - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - prompts = [ - "The capital of France is", - "The president of the United States is", - "The future of AI is", - ] - - sampling_params = SamplingParams(max_tokens=1024, temperature=0) - llm_spec.generate(prompts, sampling_params) - llm_spec.shutdown() - - -@pytest.mark.high_cuda_memory -def test_pard_rejection(): - """PARD with rejection sampling on: the block-capture rejection path - (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs - end-to-end with non-greedy sampling and produces coherent output.""" - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - spec_config = PARDDecodingConfig( - max_draft_len=4, - speculative_model=pard_model_dir, - use_rejection_sampling=True, - ) - llm_spec = LLM( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=True, - cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True), - max_batch_size=2, - kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=2048), - max_seq_len=2048, - speculative_config=spec_config, - ) - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). - sampling_params = SamplingParams( - max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234 - ) - outputs = llm_spec.generate(prompts, sampling_params) - llm_spec.shutdown() - - assert len(outputs) == len(prompts) - for out in outputs: - assert len(out.outputs[0].token_ids) > 0 - assert out.outputs[0].text.strip() - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index 979941d79e86..5f4cc74e1343 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -4,212 +4,11 @@ import torch from utils.llm_data import llm_models_root -from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.speculative.suffix_automaton import SAConfig, SuffixAutomatonManager -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SADecodingConfig - - -def get_perf_metrics(result): - """Extract performance metrics from result using built-in request_perf_metrics.""" - metrics = {} - if result.outputs and result.outputs[0].request_perf_metrics: - perf = result.outputs[0].request_perf_metrics - timing = perf.timing_metrics - # Convert timedelta to seconds - metrics["arrival_time"] = timing.arrival_time.total_seconds() - metrics["first_token_time"] = timing.first_token_time.total_seconds() - metrics["last_token_time"] = timing.last_token_time.total_seconds() - # Calculate TTFT and E2E latency - metrics["ttft"] = metrics["first_token_time"] - metrics["arrival_time"] - metrics["e2e"] = metrics["last_token_time"] - metrics["arrival_time"] - return metrics - - -# Test parameter combinations: -# - disable_overlap_scheduler: Controls scheduler mode (False=overlap enabled) -# - use_cuda_graph: Whether to use CUDA graph capture -# - attn_backend: Attention implementation (TRTLLM only - FLASHINFER not supported) -# - max_matching_ngram_size: SA matching mode (2=fixed size, -1=longest match) -# -# NOTE: FLASHINFER target decode supports multiple queries per request, but -# non-shared one-engine modes still require a separate draft KV cache. The -# draft KV metadata/manager swap is currently implemented only for TRTLLM -# attention. Shared-target-KV modes use a separate FlashInfer metadata view. -@pytest.mark.parametrize( - "disable_overlap_scheduler,use_cuda_graph,attn_backend,max_matching_ngram_size", - [ - [False, False, "TRTLLM", 2], - [False, True, "TRTLLM", 2], - [True, False, "TRTLLM", 2], - [True, True, "TRTLLM", 2], - [False, False, "TRTLLM", -1], - ], -) -@pytest.mark.high_cuda_memory -def test_llama_sa( - disable_overlap_scheduler: bool, - use_cuda_graph: bool, - attn_backend: str, - max_matching_ngram_size: int, -): - """Test SA (Suffix Automaton) speculative decoding correctness and acceptance rate. - - Verifies: - 1. Speculative decoding produces identical results to baseline - 2. SA drafting produces draft tokens that get accepted - 3. Multi-token acceptance occurs (acceptanceLength > 1) - """ - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 20: - pytest.skip("Not enough memory to load target model") - - print( - f"\nTest config: disable_overlap_scheduler={disable_overlap_scheduler}, " - f"use_cuda_graph={use_cuda_graph}, attn_backend={attn_backend}, " - f"max_matching_ngram_size={max_matching_ngram_size}" - ) - - max_batch_size = 1 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) - cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) if use_cuda_graph else None - - llm_common_config = dict( - model=llm_models_root() / "llama-3.1-model" / "Meta-Llama-3.1-8B", - backend="pytorch", - attn_backend=attn_backend, - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_num_tokens=2048, - enable_iter_perf_stats=True, - ) - - spec_config = SADecodingConfig( - max_draft_len=max_draft_len, - max_matching_ngram_size=max_matching_ngram_size, - ) - - # Use prompts that encourage repetitive patterns for better SA/ngram matching - prompts = [ - "Count from 1 to 50: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, " - "16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, " - "34, 35,", - ] - # Enable perf metrics collection via return_perf_metrics=True - sampling_params = SamplingParams( - max_tokens=64, ignore_eos=True, temperature=0, return_perf_metrics=True - ) - - # Run with speculative decoding - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - results_spec = llm_spec.generate(prompts, sampling_params) - generated_text_spec = [result.outputs[0].text for result in results_spec] - - # Get spec decoding stats before shutdown - stats = llm_spec.get_stats(timeout=5) - iterations_with_spec = [] - for stat in stats: - if "specDecodingStats" in stat: - spec_stats = stat["specDecodingStats"] - if spec_stats.get("numDraftTokens", 0) > 0: - iterations_with_spec.append(spec_stats) - - # Get perf metrics using built-in request_perf_metrics - spec_metrics = get_perf_metrics(results_spec[0]) if results_spec else {} - - llm_spec.shutdown() - - # Run reference without speculative decoding - llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) - generated_text_ref = [result.outputs[0].text for result in results_ref] +from tensorrt_llm.llmapi import SADecodingConfig - # Get perf metrics for reference - ref_metrics = get_perf_metrics(results_ref[0]) if results_ref else {} - - llm_ref.shutdown() - - # Verify 1: Identical results (correctness) - for i, (text_spec, text_ref) in enumerate(zip(generated_text_spec, generated_text_ref)): - assert text_spec == text_ref, ( - f"Prompt {i}: Spec decode result differs from baseline.\n" - f"Spec: {text_spec}\nRef: {text_ref}" - ) - print("Correctness verified: spec decode matches baseline") - - # Verify 2: Spec decoding stats show drafting occurred - assert len(iterations_with_spec) > 0, ( - f"SA should have iterations with specDecodingStats. " - f"Got {len(stats)} total stats but 0 with draft tokens." - ) - - total_draft = sum(s["numDraftTokens"] for s in iterations_with_spec) - total_accepted = sum(s["numAcceptedTokens"] for s in iterations_with_spec) - avg_acceptance_len = sum(s["acceptanceLength"] for s in iterations_with_spec) / len( - iterations_with_spec - ) - - print("Spec decoding stats:") - print(f" Iterations with drafting: {len(iterations_with_spec)}") - print(f" Total draft tokens: {total_draft}") - print(f" Total accepted tokens: {total_accepted}") - print(f" Average acceptance length: {avg_acceptance_len:.2f}") - print(f" Acceptance rate: {total_accepted / total_draft * 100:.1f}%") - - assert total_draft > 0, "SA should produce draft tokens" - assert total_accepted > 0, ( - f"SA should accept some draft tokens. " - f"Got {total_accepted} accepted out of {total_draft} drafted" - ) - - # Verify 3: Multi-token acceptance (acceptanceLength > 1) - has_multi_token_acceptance = any(s["acceptanceLength"] > 1.0 for s in iterations_with_spec) - print(f" Has multi-token acceptance: {has_multi_token_acceptance}") - - assert has_multi_token_acceptance, ( - "Expected at least one iteration with acceptanceLength > 1 for repetitive pattern" - ) - - # Print performance comparison using built-in metrics - print("\n" + "=" * 70) - print("PERFORMANCE COMPARISON (using request_perf_metrics)") - print("=" * 70) - print( - f"Config: overlap_scheduler={'enabled' if not disable_overlap_scheduler else 'disabled'}, " - f"cuda_graph={'enabled' if use_cuda_graph else 'disabled'}" - ) - print("-" * 70) - print(f"{'Metric':<30} {'Spec Decoding':<20} {'Reference':<20}") - print("-" * 70) - - # Print TTFT (Time to First Token) - ttft_spec = spec_metrics.get("ttft", None) - ttft_ref = ref_metrics.get("ttft", None) - ttft_spec_str = f"{ttft_spec * 1000:.2f} ms" if ttft_spec else "N/A" - ttft_ref_str = f"{ttft_ref * 1000:.2f} ms" if ttft_ref else "N/A" - print(f"{'TTFT':<30} {ttft_spec_str:<20} {ttft_ref_str:<20}") - - # Print E2E latency - e2e_spec = spec_metrics.get("e2e", None) - e2e_ref = ref_metrics.get("e2e", None) - e2e_spec_str = f"{e2e_spec * 1000:.2f} ms" if e2e_spec else "N/A" - e2e_ref_str = f"{e2e_ref * 1000:.2f} ms" if e2e_ref else "N/A" - print(f"{'E2E Latency':<30} {e2e_spec_str:<20} {e2e_ref_str:<20}") - - # Calculate and print speedup - if e2e_spec and e2e_ref and e2e_spec > 0: - speedup = e2e_ref / e2e_spec - print("-" * 70) - print(f"{'Speedup (E2E)':<30} {speedup:.2f}x") - print("=" * 70 + "\n") - - # Synchronize CUDA to catch any async memory errors before test completes. - # This ensures errors are attributed to this test rather than propagating - # to subsequent tests. - torch.cuda.synchronize() +sys.path.append(os.path.join(os.path.dirname(__file__), "..")) @pytest.mark.cpu_only diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py index b523ed443727..e1486a8d9696 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py @@ -1,143 +1,8 @@ import unittest -from unittest.mock import patch import pytest -import torch -from utils.llm_data import llm_models_root -from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.speculative.speculation_gate import SpeculationGate -from tensorrt_llm.llmapi import CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig -from tensorrt_llm.logger import logger - - -@pytest.fixture(scope="function") -def enforce_single_worker(monkeypatch): - """Mock functions don't work with multiple processes, so we enforce single worker.""" - monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") - yield - - -# Tests that the SpeculationGate correctly disables speculative decoding -# when the average acceptance rate drops below the threshold. -# This test uses a mock to simulate low acceptance rates and verifies -# that the spec gate triggers and disables speculation. -@pytest.mark.high_cuda_memory -def test_spec_gate_e2e(enforce_single_worker): - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - max_batch_size = 3 - max_draft_len = 4 - acceptance_rate_window_size = 3 - acceptance_rate_threshold = 0.6 - kv_cache_config = KvCacheConfig( - enable_block_reuse=False, - free_gpu_memory_fraction=0.6, - ) - cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) - - llm_common_config = dict( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=False, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - cuda_graph_config=cuda_graph_config, - enable_chunked_prefill=False, - max_num_tokens=8192, - ) - - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=True, - acceptance_rate_window_size=acceptance_rate_window_size, - acceptance_rate_threshold=acceptance_rate_threshold, - ) - - prompts = [ - "The capital of France is", - "The president of the United States is", - "What is the capital of Australia?", - ] - sampling_params = SamplingParams(max_tokens=20, temperature=0) - - # Track calls to record_acceptance_rate and the disabled state. - gate_state = {"record_calls": [], "gate_disabled": False} - - original_record_acceptance_rate = SpeculationGate.record_acceptance_rate - - def mock_record_acceptance_rate(self, acceptance_rate, sample_id=None): - """ - Mock that simulates a low true acceptance rate. - This is below the threshold of 0.6, so the gate should trigger after - the window fills. - """ - simulated_low_rate = 0.2 - disabled_now, avg = original_record_acceptance_rate(self, simulated_low_rate, sample_id) - - gate_state["record_calls"].append( - { - "original_rate": acceptance_rate, - "simulated_rate": simulated_low_rate, - "disabled_now": disabled_now, - "avg_acceptance_rate": avg, - "sample_id": sample_id, - } - ) - if disabled_now: - gate_state["gate_disabled"] = True - - return disabled_now, avg - - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - try: - with patch.object(SpeculationGate, "record_acceptance_rate", mock_record_acceptance_rate): - llm_spec.generate(prompts, sampling_params) - - # Verify the mock was called (requests completed) - assert len(gate_state["record_calls"]) > 0, "record_acceptance_rate should have been called" - - # Verify the gate was disabled after enough requests with low acceptance - assert gate_state["gate_disabled"], ( - f"Gate should have been disabled with simulated low acceptance. Calls: {gate_state['record_calls']}" - ) - - # Verify the gate triggered at the right time (after the window is filled). - # The gate should trigger on the `acceptance_rate_window_size`-th call - # (index = window - 1). - disable_indices = [ - i for i, call in enumerate(gate_state["record_calls"]) if call["disabled_now"] - ] - assert len(disable_indices) == 1, ( - f"Gate should have triggered exactly once, but triggered at indices: {disable_indices}" - ) - assert disable_indices[0] >= acceptance_rate_window_size - 1, ( - f"Gate should trigger after window ({acceptance_rate_window_size}) " - f"is filled, but triggered at index {disable_indices[0]}" - ) - - # Verify the average acceptance rate was below threshold when disabled. - disable_call = gate_state["record_calls"][disable_indices[0]] - assert disable_call["avg_acceptance_rate"] is not None - assert disable_call["avg_acceptance_rate"] < acceptance_rate_threshold, ( - f"Avg acceptance rate ({disable_call['avg_acceptance_rate']}) " - f"should be below threshold ({acceptance_rate_threshold})" - ) - - logger.debug(f"Gate correctly triggered after {disable_indices[0] + 1} requests") - logger.debug( - f"Final avg acceptance rate: {disable_call['avg_acceptance_rate']:.3f} " - f"< threshold {acceptance_rate_threshold}" - ) - finally: - llm_spec.shutdown() @pytest.mark.cpu_only diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py b/tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py deleted file mode 100644 index ebc0e6e91e62..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py +++ /dev/null @@ -1,90 +0,0 @@ -import unittest - -import pytest -import torch -from utils.llm_data import llm_models_root - -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm._torch.speculative.ngram import NGramDrafter, NGramPoolManager -from tensorrt_llm.llmapi import ( - CudaGraphConfig, - KvCacheConfig, - NGramDecodingConfig, - UserProvidedDecodingConfig, -) - - -# TODO: add disable_overlap_scheduler=False -@pytest.mark.parametrize( - "disable_overlap_scheduler,use_cuda_graph,attn_backend", - [[True, False, "TRTLLM"], [True, True, "TRTLLM"], [True, False, "FLASHINFER"]], -) -def test_llama_user_provided( - disable_overlap_scheduler: bool, use_cuda_graph: bool, attn_backend: str -): - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 20: - pytest.skip("Not enough memory to load target model") - - max_batch_size = 2 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) - cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) if use_cuda_graph else None - - llm_common_config = dict( - model=llm_models_root() / "llama-3.1-model" / "Meta-Llama-3.1-8B", - backend="pytorch", - attn_backend=attn_backend, - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_num_tokens=2048, - ) - - ngram_config = NGramDecodingConfig( - max_draft_len=max_draft_len, - max_matching_ngram_size=2, - is_keep_all=True, - is_use_oldest=True, - is_public_pool=True, - ) - - ngram_pool_manager = NGramPoolManager( - spec_config=ngram_config, - max_num_requests=max_batch_size, - ) - - drafter = NGramDrafter( - spec_config=ngram_config, - ngram_pool_manager=ngram_pool_manager, - ) - - spec_config = UserProvidedDecodingConfig( - max_draft_len=max_draft_len, - drafter=drafter, - ) - - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - sampling_params = SamplingParams(max_tokens=32) - - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - results_spec = llm_spec.generate(prompts, sampling_params) - generated_text_spec = [result.outputs[0].text for result in results_spec] - llm_spec.shutdown() - - llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) - generated_text_ref = [result.outputs[0].text for result in results_ref] - llm_ref.shutdown() - - for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): - # The spec decode algorithm currently guarantees identical results - assert text_spec == text_ref - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index d1b3c613848a..007a66428770 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -24,10 +24,8 @@ import pytest import torch -from test_common.llm_data import with_mocked_hf_download_for_single_gpu from utils.llm_data import llm_models_root -from utils.util import (skip_blackwell, skip_num_gpus_less_than, - skip_pre_blackwell) +from utils.util import skip_num_gpus_less_than, skip_pre_blackwell from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention.backends.sparse.dsa import ( @@ -35,7 +33,6 @@ from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams -from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.pyexecutor._util import ( _derive_draft_max_attention_window, _expand_attention_window_pattern_to_global_layers) @@ -49,7 +46,6 @@ INVALID_PROMPT_LOOKAHEAD_TOKEN from tensorrt_llm._torch.speculative.mtp_dynamic_tree import \ MTPEagleDynamicTreeWorker -from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig) @@ -647,280 +643,6 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): assert draft_kwargs["max_blocks"] is None -@pytest.mark.parametrize( - "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,multi_batch,attention_dp,use_hf_speculative_model", - [ - [True, "TRTLLM", True, False, False, False, False, False, False], - [False, "TRTLLM", True, False, False, False, False, False, False], - [True, "FLASHINFER", True, False, False, False, False, False, False], - [False, "FLASHINFER", True, False, False, False, False, False, False], - [False, "TRTLLM", False, True, True, False, False, False, False], - [True, "TRTLLM", False, True, True, False, False, False, False], - [True, "TRTLLM", True, False, True, True, False, False, False], - [True, "TRTLLM", True, False, True, False, False, False, False], - [True, "TRTLLM", True, False, False, True, False, False, False], - [True, "TRTLLM", False, False, False, False, False, False, False], - [False, "TRTLLM", False, False, False, False, False, False, False], - [True, "TRTLLM", False, False, False, False, True, False, False], - [True, "TRTLLM", False, False, False, False, True, True, False], - [False, "TRTLLM", False, False, False, False, True, False, False], - [True, "TRTLLM", False, False, False, True, False, False, False], - [True, "FLASHINFER", False, False, False, False, False, False, False], - [False, "FLASHINFER", False, False, False, False, False, False, False], - # Tests (mocked) speculative model auto-download from HuggingFace - [False, "TRTLLM", True, False, False, False, False, False, True], - ]) -@pytest.mark.high_cuda_memory -@with_mocked_hf_download_for_single_gpu -def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, - disable_overlap_scheduler: bool, enable_block_reuse: bool, - use_one_model: bool, enable_chunked_prefill: bool, - multi_batch: bool, attention_dp: bool, - use_hf_speculative_model: bool): - if not use_one_model: - pytest.skip("Two model Eagle3 is deprecated") - - # Eagle3 one model works with overlap scheduler and block reuse. - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - if use_hf_speculative_model: - eagle_model = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" - else: - eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - - # bs > 1 gives non-deterministic when doing IFB. There are slight chances - # that ref and spec does not match 100% - max_batch_size = 4 if multi_batch else 1 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, - max_tokens=8192) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[i for i in range(1, max_batch_size + - 1)]) if use_cuda_graph else None - - llm_common_config = dict( - model=target_model_dir, - attn_backend=attn_backend, - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - enable_attention_dp=attention_dp, - # This max_seq_len is larger than the one specified - # in the llama 3 8B eagle's config. We want to make sure - # that the draft model won't go above its max in warmup - # in this test. - max_seq_len=8192, - enable_chunked_prefill=enable_chunked_prefill, - ) - if enable_chunked_prefill: - # Use a small max_num_tokens so that the chunked prefill path gets exercised. - llm_common_config['max_num_tokens'] = 64 - - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model, - # Llama 3 does not support one model eagle. - eagle3_one_model=use_one_model, - ) - - # Create the LLM instance - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - # Acceptance rate tests - if enable_chunked_prefill: - # Use a long prompt for chunked prefill tests. - prompts = [ - "The capital of France is a city of romance, art, fashion, and cuisine. Paris is a must-visit destination for anyone who loves history, architecture, and culture. From the iconic Eiffel Tower to the world-famous Louvre Museum, Paris has something to offer for every interest and age.\nThe city is divided into 20 arrondissements, each with its own unique character and charm. The Latin Quarter is a popular area for students and young travelers, while the Champs-ƉlysĆ©es is a hub for shopping and dining. The Montmartre neighborhood is famous for its bohemian vibe and stunning views of the city.\nParis is also known for its beautiful parks and gardens, such as the Luxembourg Gardens and the Tuileries Garden. The city has a rich history, with landmarks like the Notre-Dame Cathedral and the Arc de Triomphe. Visitors can also explore the city's many museums, including the MusĆ©e d'Orsay and the MusĆ©e Rodin.\nIn addition to its cultural and historical attractions, Paris is also a great destination for foodies. The city is famous for its cuisine, including croissants, baguettes, and cheese. Visitors can sample the city's famous dishes at one of the many restaurants, cafes, and " - ] - tok_ids = [llm_spec.tokenizer.encode(prompts[0])] - else: - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - tok_ids = [llm_spec.tokenizer.encode("The future of AI is")] - if multi_batch: - # encode each prompt individually (encode(list) returns nested - # lists in transformers 5.x which prompt_inputs can't handle) - for p in prompts: - tok_ids.append(llm_spec.tokenizer.encode(p)) - - sampling_params = SamplingParams(max_tokens=128, temperature=0) - - for i in range(len(tok_ids)): - num_tokens = 0 - num_drafted = 0 - num_accepted = 0 - - for output in llm_spec.generate_async(tok_ids[i], - sampling_params, - streaming=True): - new_tokens = output.outputs[0].token_ids - num_drafted += max_draft_len - num_accepted += len(new_tokens) - num_tokens - 1 - num_tokens = len(new_tokens) - - accept_rate = num_accepted / num_drafted - assert accept_rate > 0.1 - - # Output tests - sampling_params = SamplingParams(max_tokens=10, temperature=0) - - results_spec = llm_spec.generate(prompts, sampling_params) - generated_text_spec = [result.outputs[0].text for result in results_spec] - llm_spec.shutdown() - - llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) - generated_text_ref = [result.outputs[0].text for result in results_ref] - llm_ref.shutdown() - - for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): - # The spec decode algorithm currently guarantees identical results - assert text_spec == text_ref - - -@pytest.mark.parametrize("eagle3_one_model", [True, False]) -def test_eagle3_spec_decoding_stats(eagle3_one_model): - """Test that specDecodingStats are correctly populated in metrics endpoint""" - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - # Skip if models don't exist - if not os.path.exists(target_model_dir) or not os.path.exists( - eagle_model_dir): - pytest.skip(f"Required models not found") - - kv_cache_config = KvCacheConfig(enable_block_reuse=False, - free_gpu_memory_fraction=0.6) - spec_config = Eagle3DecodingConfig( - max_draft_len=3, - speculative_model=eagle_model_dir, - eagle3_one_model=eagle3_one_model, - ) - - with LLM( - model=target_model_dir, - speculative_config=spec_config, - kv_cache_config=kv_cache_config, - disable_overlap_scheduler=not eagle3_one_model, - enable_iter_perf_stats=True, - max_batch_size=4, - ) as llm: - # Generate some output to collect stats - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - sampling_params = SamplingParams(max_tokens=20, temperature=0) - llm.generate(prompts, sampling_params) - - # Get iteration stats - stats = llm.get_stats(timeout=2) - assert len(stats) > 0, "Should have iteration stats" - - # Find iterations with speculation (generation phase) - iterations_with_spec = [] - for stat in stats: - if 'specDecodingStats' in stat: - spec_stats = stat['specDecodingStats'] - if spec_stats.get('numDraftTokens', 0) > 0: - iterations_with_spec.append(spec_stats) - - # Should have at least some iterations with spec decoding - assert len(iterations_with_spec) > 0, \ - f"Should have iterations with specDecodingStats (found {len(iterations_with_spec)})" - - # Validate specDecodingStats structure and values - for spec_stats in iterations_with_spec: - # Check all fields are present - assert 'numDraftTokens' in spec_stats - assert 'numAcceptedTokens' in spec_stats - assert 'numRequestsWithDraftTokens' in spec_stats - assert 'acceptanceLength' in spec_stats - assert 'iterLatencyMS' in spec_stats - assert 'draftOverhead' in spec_stats - - # Validate value constraints - assert spec_stats['numDraftTokens'] > 0 - assert 0 <= spec_stats['numAcceptedTokens'] <= spec_stats[ - 'numDraftTokens'] - assert spec_stats['numRequestsWithDraftTokens'] > 0 - assert spec_stats['acceptanceLength'] >= 1.0 - assert spec_stats['iterLatencyMS'] >= 0.0 - assert 0.0 <= spec_stats['draftOverhead'] <= 1.0 - - # Calculate overall acceptance rate - total_draft = sum(s['numDraftTokens'] for s in iterations_with_spec) - total_accepted = sum(s['numAcceptedTokens'] - for s in iterations_with_spec) - acceptance_rate = (total_accepted / total_draft * - 100) if total_draft > 0 else 0 - - # Should have reasonable acceptance rate for Eagle3 - assert acceptance_rate > 5.0, f"Acceptance rate too low: {acceptance_rate:.1f}%" - - -@pytest.mark.parametrize("use_cuda_graph", [True, False]) -@pytest.mark.high_cuda_memory -def test_llama_eagle3_long_prompt(use_cuda_graph): - # Eagle3 one model works with overlap scheduler and block reuse. - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - spec_config = Eagle3DecodingConfig( - max_draft_len=3, - speculative_model=eagle_model_dir, - eagle3_one_model=False, - ) - - if use_cuda_graph: - cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) - else: - cuda_graph_config = None - - llm_spec = LLM(model=target_model_dir, - speculative_config=spec_config, - max_batch_size=1, - cuda_graph_config=cuda_graph_config, - disable_overlap_scheduler=True) - - prompt = [", ".join(str(i) for i in range(1000))] - - sampling_params = SamplingParams(max_tokens=10, temperature=0) - results_spec = llm_spec.generate(prompt, sampling_params) - - generated_text_spec = [result.outputs[0].text for result in results_spec] - llm_spec.shutdown() - - llm_ref = LLM(model=target_model_dir, - max_batch_size=1, - cuda_graph_config=None, - disable_overlap_scheduler=False) - - results_ref = llm_ref.generate(prompt, sampling_params) - - generated_text_ref = [result.outputs[0].text for result in results_ref] - llm_ref.shutdown() - - # The LLM with speculation on should dynamically turn it off in this - # test since it goes beyond the max seqlen. Thus, the text should be - # _exactly_ the same, no need to use similarity scoring. - assert generated_text_spec[0] == generated_text_ref[0] - - def test_deepseek_mla_eagle3(): use_cuda_graph = True attn_backend = "TRTLLM" @@ -1029,172 +751,6 @@ def test_deepseek_mla_eagle3(): pass -@pytest.mark.parametrize("use_one_model", [True, False]) -def test_multi_eagle3(use_one_model: bool): - use_cuda_graph = True - attn_backend = "TRTLLM" - disable_overlap_scheduler = False - enable_block_reuse = False - enable_chunked_prefill = False - - # Eagle3 one model works with overlap scheduler and block reuse. - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 150: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - eagle_config = { - 'architectures': ['LlamaForCausalLMEagle3'], - 'attention_bias': False, - 'attention_dropout': 0.0, - 'bos_token_id': 128000, - 'eos_token_id': [128001, 128008, 128009], - 'eagle_config': { - 'use_aux_hidden_state': False, - 'use_input_layernorm_in_first_layer': True, - 'use_last_layernorm': True, - 'use_mtp_layernorm': False - }, - 'head_dim': 128, - 'hidden_act': 'silu', - 'hidden_size': 4096, - 'initializer_range': 0.02, - 'intermediate_size': 16384, - 'max_position_embeddings': 131072, - 'mlp_bias': False, - 'model_type': 'llama', - 'num_attention_heads': 32, - 'num_eagle_features': 1, - 'num_hidden_layers': 2, - 'num_key_value_heads': 8, - 'pretraining_tp': 1, - 'rms_norm_eps': 1e-05, - 'rope_scaling': { - 'factor': 8.0, - 'high_freq_factor': 4.0, - 'low_freq_factor': 1.0, - 'original_max_position_embeddings': 8192, - 'rope_type': 'llama3' - }, - 'rope_theta': 500000.0, - 'tie_word_embeddings': False, - 'torch_dtype': 'bfloat16', - 'transformers_version': '4.52.4', - 'use_cache': True, - 'vocab_size': 128256, - 'draft_vocab_size': 128256, - } - with tempfile.TemporaryDirectory() as temp_dir: - eagle_model_dir = Path(temp_dir) - config_path = eagle_model_dir / "config.json" - with config_path.open("w") as f: - json.dump(eagle_config, f, indent=2) - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - - # bs > 1 gives non-deterministic when doing IFB. There are slight chances - # that ref and spec does not match 100% - max_batch_size = 16 - max_draft_len = 3 - kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, - free_gpu_memory_fraction=0.5) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1]) if use_cuda_graph else None - - llm_common_config = dict( - model=target_model_dir, - attn_backend=attn_backend, - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - enable_chunked_prefill=enable_chunked_prefill, - load_format="dummy", - ) - - spec_config = Eagle3DecodingConfig(max_draft_len=max_draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, - load_format="dummy") - - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - tok_ids = llm_spec.tokenizer.encode("The future of AI is") - - sampling_params = SamplingParams(max_tokens=32, temperature=0) - for output in llm_spec.generate_async(tok_ids, - sampling_params, - streaming=True): - pass - - -@pytest.mark.parametrize("use_dynamic_tree", [False, True], - ids=["no_dynamic_tree", "dynamic_tree"]) -@pytest.mark.parametrize("use_cuda_graph", [False, True]) -@pytest.mark.high_cuda_memory -@skip_blackwell -# Opt out of MPI session reuse: the XQA JIT cubin registry is process-global -# (DecoderXQARunner::getResourceGlobal) and its lookup key does not include -# q_seq_len / is_spec_dec_tree, so running the dynamic-tree and non-dynamic-tree -# variants in one worker process launches a cubin compiled for the other -# config's q_seq_len -> CUDA_ERROR_INVALID_VALUE on Hopper. -@pytest.mark.private_mpi_session -@with_mocked_hf_download_for_single_gpu -def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, - use_cuda_graph: bool): - """Test one-model rejection sampling with and without dynamic tree.""" - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - - max_batch_size = 1 - max_draft_len = 6 - dynamic_tree_max_top_k = 10 - max_total_draft_tokens = 60 - kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1]) if use_cuda_graph else None - - llm_common_config = dict( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=True, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_seq_len=8192, - ) - - spec_config_kwargs = dict( - max_draft_len=max_draft_len, - speculative_model=eagle_model, - eagle3_one_model=True, - use_rejection_sampling=True, - ) - if use_dynamic_tree: - spec_config_kwargs.update( - use_dynamic_tree=True, - dynamic_tree_max_topK=dynamic_tree_max_top_k, - max_total_draft_tokens=max_total_draft_tokens, - ) - - llm_spec = LLM(**llm_common_config, - speculative_config=Eagle3DecodingConfig( - **spec_config_kwargs)) - - prompts = ["The president of the United States is"] - sampling_params = SamplingParams(max_tokens=20, temperature=1.0, top_p=1.0) - - results = llm_spec.generate(prompts, sampling_params) - llm_spec.shutdown() - - assert len(results) == len(prompts) - assert len(results[0].outputs[0].token_ids) > 0 - - @pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) @pytest.mark.parametrize("use_cuda_graph", [False, True]) @pytest.mark.high_cuda_memory @@ -1273,167 +829,5 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( assert text_spec == text_ref -@pytest.mark.parametrize("use_cuda_graph", [True, False]) -def test_eagle3_lora(use_cuda_graph: bool): - """Test LoRA with 3 requests and max_batch_size=4. - - This test verifies that when using LoRA modules, - the system properly applies the LoRA configurations. - """ - attn_backend = "TRTLLM" - enable_block_reuse = False - use_one_model = True - enable_chunked_prefill = False - - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - - eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - hf_lora_dir = f"{models_path}/llama-models/luotuo-lora-7b-0.1" - - # Test with 3 requests and max_batch_size=4 to trigger padding - max_batch_size = 4 - max_draft_len = 4 - kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, - max_tokens=8192) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[1, 2, 4], enable_padding=True) if use_cuda_graph else None - lora_config = LoraConfig(max_lora_rank=64, max_loras=2, max_cpu_loras=2) - - llm_common_config = dict( - model=target_model_dir, - attn_backend=attn_backend, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_seq_len=1024, - enable_chunked_prefill=enable_chunked_prefill, - lora_config=lora_config, - ) - - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model_dir, - eagle3_one_model=use_one_model, - ) - - # Create the LLM instance - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - prompts = [ - "The capital of France is", - "The president of the United States is", - "The future of AI is", - ] - lora_requests = [LoRARequest("luotuo", 1, hf_lora_dir)] * len(prompts) - - sampling_params = SamplingParams(max_tokens=20, temperature=0) - llm_spec.generate(prompts, sampling_params, lora_request=lora_requests) - llm_spec.shutdown() - - -@pytest.mark.parametrize("disable_overlap_scheduler", [False]) -@pytest.mark.parametrize("use_cuda_graph", [True]) -@pytest.mark.high_cuda_memory -@with_mocked_hf_download_for_single_gpu -def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, - disable_overlap_scheduler: bool): - """Test EAGLE3 dynamic tree speculative decoding with one-model architecture.""" - total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 - if total_mem_gb < 35: - pytest.skip("Not enough memory to load target + draft model") - - models_path = llm_models_root() - target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" - eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" - - max_batch_size = 4 - max_draft_len = 6 - dynamic_tree_max_topK = 10 - max_total_draft_tokens = 30 - kv_cache_config = KvCacheConfig(enable_block_reuse=False, - max_tokens=2048, - free_gpu_memory_fraction=0.5) - cuda_graph_config = CudaGraphConfig( - batch_sizes=[i for i in range(1, max_batch_size + - 1)]) if use_cuda_graph else None - - llm_common_config = dict( - model=target_model_dir, - attn_backend="TRTLLM", - disable_overlap_scheduler=disable_overlap_scheduler, - cuda_graph_config=cuda_graph_config, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - max_seq_len=2048, - ) - - spec_config = Eagle3DecodingConfig( - max_draft_len=max_draft_len, - speculative_model=eagle_model, - eagle3_one_model=True, - use_dynamic_tree=True, - dynamic_tree_max_topK=dynamic_tree_max_topK, - max_total_draft_tokens=max_total_draft_tokens, - ) - - # Create the LLM instance - llm_spec = LLM(**llm_common_config, speculative_config=spec_config) - - # Acceptance rate tests - prompts = [ - "The capital of France is", - "The president of the United States is", - ] - tok_ids = [llm_spec.tokenizer.encode("The future of AI is")] - - sampling_params = SamplingParams(max_tokens=128, temperature=0) - - for i in range(len(tok_ids)): - num_tokens = 0 - num_drafted = 0 - num_accepted = 0 - - for output in llm_spec.generate_async(tok_ids[i], - sampling_params, - streaming=True): - new_tokens = output.outputs[0].token_ids - num_drafted += max_draft_len - num_accepted += len(new_tokens) - num_tokens - 1 - num_tokens = len(new_tokens) - - accept_rate = num_accepted / num_drafted - assert accept_rate > 0.10 - - # Output tests: verify spec decode matches reference - sampling_params = SamplingParams(max_tokens=10, temperature=0) - - results_spec = llm_spec.generate(prompts, sampling_params) - generated_text_spec = [result.outputs[0].text for result in results_spec] - llm_spec.shutdown() - - llm_ref = LLM(**llm_common_config) - results_ref = llm_ref.generate(prompts, sampling_params) - generated_text_ref = [result.outputs[0].text for result in results_ref] - llm_ref.shutdown() - - def assert_meaningful_text(text: str) -> None: - stripped = text.strip() - assert stripped - assert "\ufffd" not in stripped - assert any(ch.isalpha() for ch in stripped) - words = stripped.lower().split() - assert not any( - len(set(words[i:i + 6])) == 1 for i in range(len(words) - 5)) - - for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): - assert_meaningful_text(text_spec) - assert_meaningful_text(text_ref) - - if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py b/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py index 9b726263e174..899c6810be85 100644 --- a/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py +++ b/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py @@ -428,15 +428,6 @@ def apply_rotary_pos_emb_ds(q, k, cos, sin, position_ids, unsqueeze_dim=1): _SMALL_MODEL_CONFIGS = { - "meta-llama/Meta-Llama-3.1-8B-Instruct": { - "model_kwargs": { - "num_hidden_layers": 1, - "hidden_size": 64, - "intermediate_size": 64, - "num_attention_heads": 2, - "num_key_value_heads": 1, - }, - }, "Qwen/Qwen3-30B-A3B": { "model_kwargs": { "num_hidden_layers": 2, @@ -571,7 +562,7 @@ def get_small_model_config(model_hub_id: str, **llm_args_kwargs) -> Dict[str, An Get the small model configuration for a given HuggingFace model hub ID. Args: - model_hub_id: The HuggingFace model hub ID (e.g., "meta-llama/Meta-Llama-3.1-8B-Instruct") + model_hub_id: The HuggingFace model hub ID (e.g., "Qwen/Qwen3-30B-A3B") Returns: Dictionary containing the model configuration diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py deleted file mode 100644 index 2517196708e2..000000000000 --- a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py +++ /dev/null @@ -1,307 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-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 signal -import subprocess -import tempfile -from contextlib import contextmanager -from pathlib import Path - -import pytest -import torch -import yaml -from _model_test_utils import get_small_model_config -from click.testing import CliRunner -from utils.cpp_paths import llm_root # noqa: F401 - -from tensorrt_llm.commands.bench import main - -# needed since LLM API uses MPI executor pool internally for TP>1, which leaks a thread on shutdown -pytestmark = pytest.mark.threadleak(enabled=False) - - -class TimeoutError(Exception): - """Exception raised when a test times out.""" - - pass - - -@contextmanager -def timeout(seconds): - """Context manager that raises TimeoutError if code block exceeds time limit. - - Args: - seconds: Maximum time in seconds to allow the code block to run - - Raises: - TimeoutError: If the code block execution exceeds the time limit - """ - - def timeout_handler(signum, frame): - raise TimeoutError(f"Test execution exceeded {seconds} seconds timeout") - - # Set the signal handler and alarm - old_handler = signal.signal(signal.SIGALRM, timeout_handler) - signal.alarm(seconds) - try: - yield - finally: - # Restore the old signal handler and cancel the alarm - signal.alarm(0) - signal.signal(signal.SIGALRM, old_handler) - - -@pytest.fixture(scope="module", autouse=True) -def prewarm_flashinfer_jit(): - """Pre-warm FlashInfer JIT kernels before multi-GPU tests. - - This prevents a race condition where multiple MPI ranks try to JIT-compile - FlashInfer kernels simultaneously to the same cache directory, causing - Ninja build failures like: "ninja: error: opening build log: No such file or directory" - - By triggering the compilation in the main process first, the kernels are - cached and available for all worker ranks. - """ - try: - import flashinfer - import flashinfer.page - import flashinfer.sampling - - if torch.cuda.is_available(): - # Prevent concurrent JIT warmup across multiple pytest processes (e.g., xdist). - try: - import fcntl # Linux-only - except ImportError: - fcntl = None - - lock_f = None - if fcntl is not None: - import pathlib - import tempfile - - lock_path = pathlib.Path(tempfile.gettempdir()) / "flashinfer_jit_prewarm.lock" - lock_f = open(lock_path, "w") - fcntl.flock(lock_f.fileno(), fcntl.LOCK_EX) - # Create dummy tensors to trigger kernel JIT compilation - with torch.no_grad(): - device = torch.device("cuda:0") - - # Trigger page kernel compilation - try: - # Force module loading (this triggers JIT compilation) - _ = flashinfer.page.gen_page_module() - except Exception as exc: # noqa: BLE001 - import warnings - - warnings.warn(f"FlashInfer page-kernel prewarm failed: {exc!r}", RuntimeWarning) - - # Trigger sampling kernel compilation - try: - dummy_probs = torch.softmax(torch.randn(1, 100, device=device), dim=-1) - _ = flashinfer.sampling.sampling_from_probs(dummy_probs, deterministic=True) - except Exception as exc: # noqa: BLE001 - import warnings - - warnings.warn( - f"FlashInfer sampling-kernel prewarm failed: {exc!r}", RuntimeWarning - ) - - torch.cuda.empty_cache() - if lock_f is not None: - lock_f.close() - - except ImportError: - pass # FlashInfer not available - - yield - - -@pytest.fixture(scope="module") -def shared_dataset(llm_root): # noqa: F811 - """Prepare dataset once for all tests in this module.""" - model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct" - config = get_small_model_config(model_name) - with tempfile.TemporaryDirectory() as temp_dir: - dataset_path = _prepare_dataset( - llm_root, temp_dir, config["args"]["model"], num_requests=10 - ) - # Read dataset content to return it (temp_dir will be deleted) - with open(dataset_path, "r") as f: - dataset_content = f.read() - yield dataset_content - - -def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_requests: int = 10): - """Prepare a synthetic dataset for benchmarking.""" - _DATASET_NAME = "synthetic_128_128.txt" - dataset_path = Path(temp_dir, _DATASET_NAME) - - # Generate a small dataset to run a test - matching workload configuration. - # Write straight to a file via --output (the documented usage) instead of - # scraping --stdout: trtllm-bench prints an import-time banner on stdout that - # would otherwise corrupt the captured JSONL. - command = [ - "trtllm-bench", - "--model", - model_path_or_name, - "prepare-dataset", - "--output", - str(dataset_path), - "token-norm-dist", - "--input-mean", - "128", - "--output-mean", - "128", - "--input-stdev", - "0", - "--output-stdev", - "0", - "--num-requests", - str(num_requests), - ] - print(f"Running command: {' '.join(command)}") - result = subprocess.run(command, cwd=str(temp_dir), capture_output=True, text=True, timeout=300) - if result.returncode != 0: - raise RuntimeError(f"Failed to prepare dataset: {result.stderr}") - return dataset_path - - -@pytest.mark.parametrize( - "allreduce_strategy", - [ - "AUTO", - "ONESHOT", - "TWOSHOT", - "MIN_LATENCY", - "NCCL", - "SYMM_MEM", - ], -) -def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy): # noqa: F811 - """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs. - - Configuration: - The allreduce_strategy is set in the transforms config: - ```yaml - transforms: - detect_sharding: - allreduce_strategy: "ONESHOT" # or AUTO, NCCL, TWOSHOT, etc. - ``` - - Test configuration: - - Model: Llama-3.1-8B with TP=2 - - Dataset: 10 synthetic requests (128 input, 128 output tokens) - - Timeout: 300 seconds to catch hangs - - Skipped if fewer than 2 GPUs available - - Args: - llm_root: Root directory fixture - shared_dataset: Shared dataset fixture (prepared once for all test runs) - allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL) - """ - # Fixed timeout for all strategies (5 minutes should be enough) - TEST_TIMEOUT_SECONDS = 300 - - model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct" - # Override hidden_size to a multiple of one warp's worth of 128-bit accesses (32 threads * - # 8 fp16 elements/access = 256) so the fused all-reduce/RMSNorm kernel never launches a - # partial-warp block. This isolates whether CI failures are specific to the partial-warp - # code path or are unrelated infra flakiness that also affects the full-warp path. - config = get_small_model_config( - model_name, - model_kwargs={ - "num_hidden_layers": 1, - "hidden_size": 256, - "intermediate_size": 256, - "num_attention_heads": 2, - "num_key_value_heads": 1, - }, - ) - tp_size = 2 - max_batch_size = 256 - max_num_tokens = 8192 - - if not torch.cuda.is_available() or torch.cuda.device_count() < tp_size: - pytest.skip(f"Allreduce strategy test requires at least {tp_size} GPUs, skipping") - - with tempfile.TemporaryDirectory() as temp_dir: - # Write shared dataset to temp location - dataset_path = Path(temp_dir, "synthetic_128_128.txt") - with open(dataset_path, "w") as f: - f.write(shared_dataset) - - # Create configuration with specified allreduce strategy in transforms - extra_llm_api_options_path = f"{temp_dir}/extra_llm_api_options.yaml" - with open(extra_llm_api_options_path, "w") as f: - yaml.dump( - { - **config["args"], - "max_batch_size": max_batch_size, - "max_num_tokens": max_num_tokens, - "max_seq_len": 256, - "transforms": { - "detect_sharding": { - "stage": "sharding", - "allreduce_strategy": allreduce_strategy, - }, - "compile_model": { - "stage": "compile", - "backend": "torch-cudagraph", - "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128, 256], - }, - }, - }, - f, - ) - - # Run benchmark with specified allreduce strategy with timeout protection - runner = CliRunner() - args = [ - "--model", - model_name, - ] - - # Only pass --model_path if it's a local filesystem path - # Note: --model_path must come BEFORE the subcommand (throughput) - if str(config["args"]["model"]).startswith("/"): - args.extend(["--model_path", str(config["args"]["model"])]) - - # Add the subcommand and its options - args.extend( - [ - "throughput", - "--backend", - "_autodeploy", - "--dataset", - str(dataset_path), - "--extra_llm_api_options", - extra_llm_api_options_path, - "--tp", - str(tp_size), - "--max_batch_size", - str(max_batch_size), - "--max_num_tokens", - str(max_num_tokens), - ] - ) - - try: - with timeout(TEST_TIMEOUT_SECONDS): - result = runner.invoke(main, args, catch_exceptions=False) - assert result.exit_code == 0, f"Benchmark failed with output: {result.output}" - except TimeoutError as e: - pytest.fail( - f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. " - f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}" - ) diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py deleted file mode 100644 index 08c449198474..000000000000 --- a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-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. -"""Testing build_and_run_ad end2end.""" - -import pytest - -__extra_import_path__ = ["~/examples/auto_deploy"] -from _model_test_utils import get_small_model_config -from build_and_run_ad import ExperimentConfig, main - - -@pytest.mark.parametrize("world_size", [1, 2]) -@pytest.mark.parametrize( - "model_hub_id, llm_extra_args", - [ - ( - "meta-llama/Meta-Llama-3.1-8B-Instruct", - { - "transforms": { - "insert_cached_attention": {"backend": "flashinfer"}, - # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/9878 - # "compile_model": {"backend": "torch-opt"}, - "compile_model": {"backend": "torch-cudagraph"}, - }, - }, - ), - ], -) -def test_build_ad(world_size: int, model_hub_id: str, llm_extra_args: dict): - experiment_config = get_small_model_config(model_hub_id, **llm_extra_args) - - experiment_config["args"]["world_size"] = world_size - experiment_config["args"]["runtime"] = "trtllm" # Default runtime set to trtllm - - experiment_config = ExperimentConfig(**experiment_config) - print(f"Experiment Config: {experiment_config}") - main(experiment_config) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py index a146136bf820..8bb0500803a0 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py @@ -16,7 +16,6 @@ """Unit tests for Eagle3 model with AutoDeploy.""" from pathlib import Path -from typing import Any, ClassVar, Dict from unittest.mock import patch import pytest @@ -25,15 +24,8 @@ import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm -from tensorrt_llm._torch.auto_deploy.models.custom.modeling_eagle import ( - Eagle3DraftOutput, - EagleConfig, - EagleDrafterForCausalLM, - EagleRMSNorm, - EagleWrapper, -) +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_eagle import EagleRMSNorm, EagleWrapper from tensorrt_llm._torch.auto_deploy.models.eagle import EagleDrafterFactory -from tensorrt_llm._torch.auto_deploy.models.factory import ModelFactoryRegistry from tensorrt_llm._torch.auto_deploy.utils.node_utils import ( get_weight_shape, infer_draft_embedding_size, @@ -68,98 +60,6 @@ "topk_group": 1, } -############################################################################### -# Mock classes for standalone Eagle testing -# -# These classes enable unit testing the Eagle checkpoint without a target model. -# In production speculative decoding, real hidden states come from the target model. -# For testing, MockEagle3ModelForCausalLM generates random hidden states. -############################################################################### - - -class MockEagleConfig(EagleConfig): - """Config for standalone Eagle testing with embedding/lm_head loaded from checkpoint. - - In production, Eagle shares embedding/lm_head with the target model. - For standalone testing, we need to load these from the checkpoint. - """ - - _drafter_defaults: ClassVar[Dict[str, Dict[str, Any]]] = { - "llama": { - "load_embedding_from_target": False, - "load_lm_head_from_target": False, - "num_capture_layers": 1, - }, - } - - -class MockEagle3ModelForCausalLM(EagleDrafterForCausalLM): - """Test wrapper that provides random hidden states for standalone Eagle testing. - - In production speculative decoding, real hidden states come from the target model. - This mock class generates random hidden states for testing the Eagle model in isolation. - """ - - def __init__(self, config): - super().__init__(config) - self._hidden_size = config.hidden_size - self._dtype = config.dtype - - def forward(self, input_ids, position_ids, input_embeds=None, **kwargs): - assert self.model.embed_tokens is not None, ( - "embed_tokens must be set before running standalone Eagle model." - ) - assert self.lm_head is not None, ( - "lm_head must be set before running standalone Eagle model." - ) - - if input_embeds is None: - inputs_embeds = self.model.embed_tokens(input_ids) - - # Inject mock hidden states if not provided - if "hidden_states" not in kwargs: - batch_size, seq_len = input_ids.shape - kwargs["hidden_states"] = torch.randn( - (batch_size, seq_len, self._hidden_size), - dtype=self._dtype, - device=input_ids.device, - ) - draft_output = super().forward(inputs_embeds, position_ids, **kwargs) - logits = self.lm_head(draft_output.norm_hidden_state) - return Eagle3DraftOutput(logits=logits, last_hidden_state=draft_output.last_hidden_state) - - -class MockEagleDrafterFactory(EagleDrafterFactory): - """Test factory that uses MockEagle3ModelForCausalLM for standalone Eagle testing. - - This factory directly builds MockEagle3ModelForCausalLM with MockEagleConfig, - which loads embedding/lm_head from checkpoint for standalone testing. - """ - - def _build_model(self, device): - from contextlib import nullcontext - - from accelerate import init_empty_weights - - model_config, unused_kwargs = self._get_model_config() - # transformers>=5.5 applies @dataclass(kw_only=True) to PretrainedConfig - # subclasses, overriding EagleConfig.__init__. Use the factory classmethod. - model_config = MockEagleConfig.from_base_config(model_config, model_config.model_type) - - with (init_empty_weights if device == "meta" else nullcontext)(): - model = MockEagle3ModelForCausalLM._from_config(model_config, **unused_kwargs) - - if device == "meta": - if hasattr(model, "post_init"): - model.post_init() - else: - model.to(device) - - self._checkpoint_conversion_mapping = getattr(model, "_checkpoint_conversion_mapping", None) - model.eval() - - return model - def _build_small_draft_factory( model_hub_id: str, model_kwargs: dict | None = None @@ -185,48 +85,6 @@ def test_eagle_rmsnorm_keeps_fp32_weights(): assert norm.weight.dtype == torch.float32 -@pytest.fixture -def register_mock_eagle_factory(): - """Register MockEagleDrafterFactory for the test and clean up afterwards. - - This fixture temporarily registers the mock factory with ModelFactoryRegistry, - allowing tests to use model_factory="MockEagleDrafter", and removes the - registration after the test completes. - """ - ModelFactoryRegistry._registry["MockEagleDrafter"] = MockEagleDrafterFactory - yield - ModelFactoryRegistry._registry.pop("MockEagleDrafter", None) - - -def test_build_ad_eagle(register_mock_eagle_factory): - """Test building Eagle model with AutoDeploy using MockEagleDrafterFactory. - - This test uses the MockEagleDrafterFactory which builds MockEagle3ModelForCausalLM, - a mock model that generates random hidden states for standalone Eagle testing. - """ - llm_extra_args = { - "model_factory": "MockEagleDrafter", - "transforms": { - "insert_cached_attention": {"backend": "trtllm"}, - "compile_model": { - "backend": "torch-simple", - "piecewise_enabled": False, - }, - }, - } - experiment_config = get_small_model_config(EAGLE_MODEL_HUB_ID, **llm_extra_args) - experiment_config["args"]["runtime"] = "demollm" - experiment_config["args"]["world_size"] = 0 - experiment_config["args"]["tokenizer"] = hf_id_to_local_model_dir( - "meta-llama/Meta-Llama-3.1-8B-Instruct" - ) - - print(f"Experiment Config: {experiment_config}") - experiment_config = ExperimentConfig(**experiment_config) - - main(experiment_config) - - def test_eagle_model_torch_export(): """Test that Eagle3Model can be exported with torch.export. diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py index b05dc5949634..bcd24492a53f 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py @@ -84,32 +84,6 @@ def _check_ad_config(experiment_config: ExperimentConfig, llm_args: LlmArgs): @pytest.mark.parametrize( "model_hub_id, llm_extra_args", [ - ( - "meta-llama/Meta-Llama-3.1-8B-Instruct", - { - "kv_cache_config": { - "free_gpu_memory_fraction": 0.0001, - }, - "transforms": { - "insert_cached_attention": {"backend": "flashinfer"}, - # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/9878 - # "compile_model": {"backend": "torch-opt"}, - "compile_model": { - "backend": "torch-cudagraph", - "cuda_graph_batch_sizes": [1, 2], - }, - }, - }, - ), - ( - "meta-llama/Meta-Llama-3.1-8B-Instruct", - { - "transforms": { - "transformers_replace_cached_attn": {"backend": "flashinfer"}, - }, - "mode": "transformers", - }, - ), ( "Qwen/Qwen3-30B-A3B", { diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py deleted file mode 100644 index 2ee4ce344b0e..000000000000 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 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. - -from _model_test_utils import get_small_model_config - -from tensorrt_llm.llmapi import GuidedDecodingParams - -__extra_import_path__ = ["~/examples/auto_deploy"] -from build_and_run_ad import ExperimentConfig, main - - -def test_ad_guided_decoding_regex_e2e(): - """Test guided decoding with regex pattern validation using the build_and_run_ad main().""" - test_case = { - "prompt": "What is the capital of France?", - "regex": r"I don't know, I am a randomly initialized model|Paris", - "valid_responses": ["I don't know, I am a randomly initialized model", "Paris"], - } - - guided_decoding_backend = "xgrammar" - - experiment_config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") - - # DemoLLM runtime does not support guided decoding. Need to set runtime to trtllm. - experiment_config["args"]["runtime"] = "trtllm" - experiment_config["args"]["world_size"] = 1 - # NOTE: trtllm attention backend fails on B200 (likely illegal memory access); use flashinfer. - experiment_config["args"]["attn_backend"] = "flashinfer" - experiment_config["args"]["guided_decoding_backend"] = guided_decoding_backend - - experiment_config["prompt"]["batch_size"] = 1 - experiment_config["prompt"]["queries"] = test_case["prompt"] - - cfg = ExperimentConfig(**experiment_config) - - # Need to introduce the guided decoding params after ExperimentConfig construction - # because otherwise they get unpacked as a dict. - cfg.prompt.sp_kwargs = { - "max_tokens": 10, - "top_k": None, - "temperature": 0.1, - "guided_decoding": GuidedDecodingParams(regex=test_case["regex"]), - } - - print(f"Experiment config: {experiment_config}") - print("Generating outputs...") - results = main(cfg) - print("Results:", results) - - # Parse and validate: output should be a prefix of one of the valid responses - prompts_and_outputs = results["prompts_and_outputs"] - assert len(prompts_and_outputs) == 1 - generated_text = prompts_and_outputs[0][1].strip() - - valid_responses = test_case["valid_responses"] - is_valid_prefix = any(response.startswith(generated_text) for response in valid_responses) - assert is_valid_prefix, ( - f"Output is not a valid prefix of any expected response.\n" - f"Generated: '{generated_text}'\n" - f"Valid responses: {valid_responses}" - ) diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 849e7e8eefa9..3d287e3db6df 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -14,41 +14,9 @@ # limitations under the License. -import torch from _model_test_utils import get_small_model_config -from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm -from tensorrt_llm._torch.auto_deploy.models.eagle import EagleOneModelFactory -from tensorrt_llm._torch.auto_deploy.transform.interface import TransformConfig -from tensorrt_llm._torch.auto_deploy.transform.library.hidden_states import ( - DetectHiddenStatesForCapture, -) -from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens -from tensorrt_llm.llmapi import Eagle3DecodingConfig, MTPDecodingConfig - -__extra_import_path__ = ["~/examples/auto_deploy"] -from build_and_run_ad import ExperimentConfig, main -from test_common.llm_data import hf_id_to_local_model_dir - - -def get_extra_seq_len_for_kv_cache(llm_args) -> int: - """Mirror the current extra-KV sizing logic used by the runtime.""" - extra = 0 - spec_config = llm_args.speculative_config - if not llm_args.disable_overlap_scheduler: - extra += 1 - if spec_config is not None: - extra += spec_config.tokens_per_gen_step - 1 - - if spec_config is not None: - extra += spec_config.tokens_per_gen_step - 1 - extra += get_num_extra_kv_tokens(spec_config) - - return extra - - -def piecewise_disabled_transforms(): - return {"compile_model": {"piecewise_enabled": False}} +from tensorrt_llm.llmapi import MTPDecodingConfig def test_super_mtp_smoke(): @@ -166,129 +134,3 @@ def test_super_mtp_ssm_replay_smoke(): prompts_and_outputs = results["prompts_and_outputs"] assert len(prompts_and_outputs) == 1 - - -def test_kv_cache_extra_seq_len_for_spec_dec(): - """Test that get_extra_seq_len_for_kv_cache computes correct extra capacity.""" - from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs - - # Case 1: No spec config, no overlap - args_no_spec = LlmArgs( - model="meta-llama/Meta-Llama-3.1-8B-Instruct", - disable_overlap_scheduler=True, - ) - assert get_extra_seq_len_for_kv_cache(args_no_spec) == 0 - - # Case 2: No spec config, with overlap - args_overlap = LlmArgs( - model="meta-llama/Meta-Llama-3.1-8B-Instruct", - disable_overlap_scheduler=False, - ) - assert get_extra_seq_len_for_kv_cache(args_overlap) == 1 # overlap adds +1 - - # Case 3: Eagle3 one-model, overlap disabled - spec_config = Eagle3DecodingConfig( - max_draft_len=3, - speculative_model="some/model", - eagle3_one_model=True, - ) - args_eagle = LlmArgs( - model="meta-llama/Meta-Llama-3.1-8B-Instruct", - speculative_config=spec_config, - disable_overlap_scheduler=True, - transforms=piecewise_disabled_transforms(), - ) - extra = get_extra_seq_len_for_kv_cache(args_eagle) - # Should include max_total_draft_tokens + get_num_extra_kv_tokens (max_draft_len - 1) - assert extra > 0 - assert extra == spec_config.max_total_draft_tokens + (spec_config.max_draft_len - 1) - - # Case 4: Eagle3 one-model, overlap enabled - args_eagle_overlap = LlmArgs( - model="meta-llama/Meta-Llama-3.1-8B-Instruct", - speculative_config=spec_config, - disable_overlap_scheduler=False, - transforms=piecewise_disabled_transforms(), - ) - extra_overlap = get_extra_seq_len_for_kv_cache(args_eagle_overlap) - # Should be more than without overlap - assert extra_overlap > extra - - -def test_mtp_autodeploy_uses_eagle_one_model_capture(): - from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs - - model = "meta-llama/Meta-Llama-3.1-8B-Instruct" - args = LlmArgs( - model=model, - speculative_config=MTPDecodingConfig( - num_nextn_predict_layers=3, - mtp_eagle_one_model=True, - ), - transforms=piecewise_disabled_transforms(), - ) - - assert isinstance(args.speculative_config, MTPDecodingConfig) - assert args.model_factory == "eagle_one_model" - assert args.transforms["detect_hidden_states_for_capture"]["enabled"] is True - assert args.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] == {-1} - - -def test_detect_hidden_states_capture_last_layer_for_mtp_eagle_one_model(): - from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs - - config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") - config["args"].setdefault("transforms", {}).setdefault("compile_model", {})[ - "piecewise_enabled" - ] = False - - args = LlmArgs( - **config["args"], - speculative_config=MTPDecodingConfig( - num_nextn_predict_layers=3, - mtp_eagle_one_model=True, - speculative_model=config["args"]["model"], - ), - ) - - factory = args.create_factory() - assert isinstance(factory, EagleOneModelFactory) - - model = factory.target_factory.build_model("meta") - input_ids = torch.ones((1, 8), dtype=torch.int64) - position_ids = torch.arange(8, dtype=torch.int64).unsqueeze(0) - gm = torch_export_to_gm( - model, - args=(input_ids, position_ids), - ) - - transform = DetectHiddenStatesForCapture( - config=TransformConfig( - stage="pattern_matcher", - eagle3_layers_to_capture={-1}, - ) - ) - - original_residual_nodes = transform.collect_residual_add_nodes(gm) - assert original_residual_nodes - last_layer = max(original_residual_nodes) - last_layer_residual = original_residual_nodes[last_layer] - expected_arg_names = tuple( - arg.name if isinstance(arg, torch.fx.Node) else arg for arg in last_layer_residual.args - ) - - gm, info = transform._apply(gm, None, None, None) - - capture_nodes = [ - node - for node in gm.graph.nodes - if node.op == "call_function" - and node.target == torch.ops.auto_deploy.residual_add_for_capture.default - ] - - assert info.num_matches == 1 - assert len(capture_nodes) == 1 - capture_arg_names = tuple( - arg.name if isinstance(arg, torch.fx.Node) else arg for arg in capture_nodes[0].args - ) - assert capture_arg_names == expected_arg_names diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py deleted file mode 100644 index 01456c01712e..000000000000 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py +++ /dev/null @@ -1,117 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-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 multiprocessing as mp -import time - -import pytest -import yaml -from _model_test_utils import get_small_model_config # type: ignore -from click.testing import CliRunner -from openai import OpenAI - -from tensorrt_llm._utils import get_free_port -from tensorrt_llm.commands.serve import main as serve_main - - -def _run_serve_with_click(args): - runner = CliRunner() - # Blocks while server runs - result = runner.invoke(serve_main, args, catch_exceptions=False) - if result.exit_code != 0: - raise SystemExit(result.exit_code) - - -@pytest.mark.timeout(500) -def test_trtllm_serve_openai_chat_completion(tmp_path): - # Prepare small model config and extra options yaml - config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") - extra_args = config["args"] - - # NOTE: trtllm attention backend fails on B200 (likely illegal memory access); use flashinfer. - extra_args["attn_backend"] = "flashinfer" - extra_options_path = tmp_path / "extra_llm_api_options.yaml" - with open(extra_options_path, "w") as f: - yaml.safe_dump(extra_args, f) - - host = "127.0.0.1" - port = get_free_port() - - # Use the same `model` string for server and client requests - model_id = extra_args["model"] - - args = [ - "serve", - f"{model_id}", - "--backend", - "_autodeploy", - "--host", - host, - "--port", - str(port), - "--extra_llm_api_options", - str(extra_options_path), - ] - - ctx = mp.get_context("spawn") - server = ctx.Process(target=_run_serve_with_click, args=(args,)) - server.start() - - try: - # Wait for server to be ready by polling /v1/models via OpenAI client - client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key="tensorrt_llm") - - start_time = time.time() - last_err = None - while time.time() - start_time < 300: - if not server.is_alive(): - raise RuntimeError("Server process exited prematurely") - try: - # Lightweight readiness probe - _ = client.models.list() - break - except Exception as e: # noqa: BLE001 - last_err = e - time.sleep(1) - else: - raise TimeoutError(f"Server did not become ready in time: {last_err}") - - # Send a small chat completion request - resp = client.chat.completions.create( - model=model_id, - messages=[ - {"role": "system", "content": "you are a helpful assistant"}, - {"role": "user", "content": "Say 'ok'"}, - ], - max_tokens=8, - ) - - # print response - print(f"{resp=}") - - assert hasattr(resp, "choices") and len(resp.choices) > 0 - first = resp.choices[0] - # new OpenAI client returns .message for chat completions - assert getattr(first, "message", None) is not None - # Content may be a string or a structured list depending on client version - _ = getattr(first.message, "content", None) - - finally: - # Terminate server and clean up - if server.is_alive(): - server.terminate() - server.join(timeout=20) - if server.is_alive(): - server.kill() - server.join(timeout=20) diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py index 10b802fe7173..490bcfd26570 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py @@ -22,10 +22,7 @@ from tensorrt_llm import DisaggregatedParams, SamplingParams from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM -from tensorrt_llm.llmapi import Eagle3DecodingConfig -LLAMA_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct" -EAGLE3_MODEL_ID = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" DEEPSEEK_MODEL_ID = "deepseek-ai/DeepSeek-V3" DEEPSEEK_DISAGG_TRANSFORMS = { "insert_cached_attention": {"backend": "triton"}, @@ -100,11 +97,6 @@ def create_generation_params(context_output): return replace(params, request_type="generation_only") -def has_draft_tokens(output): - params = output.disaggregated_params - return params is not None and params.draft_tokens is not None and len(params.draft_tokens) > 0 - - def run_live_disagg_smoke( model_id, attn_backend, @@ -212,16 +204,6 @@ def run_live_batch_disagg_smoke(model_id, attn_backend, compile_backend, config_ GENERIC_DISAGG_SMOKE_CASES = [ - pytest.param(LLAMA_MODEL_ID, "trtllm", "torch-simple", {}, id="llama-trtllm-simple"), - pytest.param(LLAMA_MODEL_ID, "trtllm", "torch-cudagraph", {}, id="llama-trtllm-cudagraph"), - pytest.param(LLAMA_MODEL_ID, "flashinfer", "torch-simple", {}, id="llama-flashinfer-simple"), - pytest.param( - LLAMA_MODEL_ID, - "flashinfer", - "torch-cudagraph", - {}, - id="llama-flashinfer-cudagraph", - ), pytest.param( DEEPSEEK_MODEL_ID, "trtllm", @@ -261,44 +243,3 @@ def test_autodeploy_disaggregated_batch_smoke( ) run_live_batch_disagg_smoke(model_id, attn_backend, compile_backend, config_overrides) - - -def test_autodeploy_disaggregated_eagle3_smoke(): - target_model_config = get_small_model_config(LLAMA_MODEL_ID) - eagle3_model_config = get_small_model_config(EAGLE3_MODEL_ID) - target_model_kwargs = { - **target_model_config["args"]["model_kwargs"], - "num_hidden_layers": 3, - } - speculative_config = Eagle3DecodingConfig( - max_draft_len=3, - speculative_model=eagle3_model_config["args"]["model"], - eagle3_one_model=True, - eagle3_layers_to_capture={0, 1, 2}, - ) - speculative_model_kwargs = { - **target_model_kwargs, - **eagle3_model_config["args"]["model_kwargs"], - "torch_dtype": "bfloat16", - } - - # This is intentionally a smoke test: small_model_config_disagg uses - # skip_loading_weights=True, so the meaningful assertions are that one-model - # Eagle builds with a reduced target/draft pair and carries draft-token - # metadata through the live disaggregated handoff. Force the draft dtype to - # match the BF16 Llama target because shared KV cache management requires - # target and draft KV resources to have the same dtype. Use three reduced - # target layers to match Llama Eagle3's default three-layer capture. - # Weighted acceptance and quality coverage belong in integration tests. - context_output, generation_output = run_live_disagg_smoke( - LLAMA_MODEL_ID, - "flashinfer", - "torch-simple", - common_config_overrides={ - "model_kwargs": target_model_kwargs, - "speculative_config": speculative_config, - "speculative_model_kwargs": speculative_model_kwargs, - }, - ) - assert has_draft_tokens(context_output) - assert has_draft_tokens(generation_output) diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py deleted file mode 100644 index 02b2b4421b1b..000000000000 --- a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py +++ /dev/null @@ -1,213 +0,0 @@ -import os -import platform -import socket -import time - -import openai -import pytest -from test_common.http_utils import (wait_for_endpoint_down, - wait_for_endpoint_ready) - -from tensorrt_llm._utils import get_sm_version - -from ..test_llm import get_model_path -from .openai_server import RemoteDisaggOpenAIServer, RemoteOpenAIServer -from .utils import (expand_slurm_nodelist, wait_for_endpoint_down, - wait_for_endpoint_ready) - -RANK = int(os.environ.get("SLURM_PROCID", 0)) -NODE_RANK = int(os.environ.get("SLURM_NODEID", 0)) -NODE_LIST = expand_slurm_nodelist(os.environ.get("SLURM_NODELIST", "")) -SLURM_NTASKS_PER_NODE = int(os.environ.get("SLURM_NTASKS_PER_NODE", 1)) - -pytestmark = pytest.mark.threadleak(enabled=False) - -# This test assumes that there are >2 nodes, we run ctx/disagg-server/client on the first node, -# and run gen the second node. -# This is a multi-node test, and will not be scheduled to the same node running other tests -# using fixed ports should be safe. -CTX_SERVER_PORT = 8001 -GEN_SERVER_PORT = 8002 -DISAGG_SERVER_PORT = 8000 - - -# Exclude the current node from the node list, then return other nodes by idx -def get_the_other_host(idx=0): - assert len(NODE_LIST) >= 2 - node_list = NODE_LIST.copy() - curr_host = socket.gethostname() - if curr_host in NODE_LIST: - # gethostname returns the exact node name in node list - node_list.remove(curr_host) - else: - # gethostname returns the full domain - curr_host = curr_host.split('.')[0] - assert curr_host in node_list - node_list.remove(curr_host) - return node_list[idx] - - -def is_ctx_node(): - return NODE_RANK == 0 - - -def is_gen_node(): - return NODE_RANK == 1 - - -def is_disagg_node(): - return NODE_RANK == 0 - - -# The test is run on multinodes but only the first node's output is used for assertion -def is_pytest_node(): - return NODE_RANK == 0 - - -def env(): - # Remove MPI related environment variables to isolate the ctx/gen processes - # so that they will not be in the same MPI communicator, otherwise the rank and world_size may mismatch - e = { - k: v - for k, v in os.environ.items() - if not ('PMI_' in k or 'OMPI_' in k or 'PMIX_' in k or 'SLURM_' in k) - and k not in ["UCX_TLS", "UCX_NET_DEVICES"] # avoid UCX failure on oci - } - # Some GB300 machines have NICs that misbehave with UCX's default transport - # auto-selection, so UCX_TLS must be set explicitly there. Identify GB300 via - # sm_103 (Blackwell Ultra) + aarch64 (Grace) -- this excludes HGX B300 - # (also sm_103 but x86_64) and avoids depending on the GPU device name string. - if get_sm_version() == 103 and platform.machine().lower() == "aarch64": - e["UCX_TLS"] = "cuda_copy,cuda_ipc,sm,self,tcp" - return e - - -@pytest.fixture(scope="module") -def model_name(): - return "llama-3.1-model/Llama-3.1-8B-Instruct" - - -@pytest.fixture(scope="module", params=['pytorch'], ids=["pytorch"]) -def backend(request): - return request.param - - -@pytest.fixture( - scope="module", - params=[(1, 1), (2, 1), (1, 2)], - ids=lambda tp_pp_size: f'ctx_tp{tp_pp_size[0]}pp{tp_pp_size[1]}') -def ctx_tp_pp_size(request): - return request.param - - -@pytest.fixture( - scope="module", - params=[(1, 1), (2, 1), (1, 2)], - ids=lambda tp_pp_size: f'gen_tp{tp_pp_size[0]}pp{tp_pp_size[1]}') -def gen_tp_pp_size(request): - return request.param - - -@pytest.fixture(scope="module") -def worker(model_name: str, ctx_tp_pp_size: tuple, gen_tp_pp_size: tuple): - extra_config = { - "cache_transceiver_config": { - "backend": "DEFAULT" - }, - "kv_cache_config": { - "free_gpu_memory_fraction": 0.5, - "enable_block_reuse": False, - }, - "disable_overlap_scheduler": True, - } - if is_ctx_node(): - print(f"starting ctx_server for rank {RANK} node rank {NODE_RANK}") - model_path = get_model_path(model_name) - tp_size, pp_size = ctx_tp_pp_size - args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] - with RemoteOpenAIServer(model_path, - port=CTX_SERVER_PORT, - cli_args=args, - host="0.0.0.0", - env=env(), - llmapi_launch=False, - rank=RANK % SLURM_NTASKS_PER_NODE, - extra_config=extra_config) as server: - yield server - elif is_gen_node(): - print(f"starting gen_server for rank {RANK} node rank {NODE_RANK}") - model_path = get_model_path(model_name) - tp_size, pp_size = gen_tp_pp_size - args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] - with RemoteOpenAIServer(model_path, - port=GEN_SERVER_PORT, - cli_args=args, - host="0.0.0.0", - env=env(), - llmapi_launch=False, - rank=RANK % SLURM_NTASKS_PER_NODE, - extra_config=extra_config) as server: - yield server - else: - yield None - - -@pytest.fixture(scope="module") -def disagg_server(worker: RemoteOpenAIServer): - if is_disagg_node(): - print(f"starting disagg_server for rank {RANK} node rank {NODE_RANK}") - ctx_url = f"localhost:8001" # Use localhost since the ctx server is on the same node - # TODO: Hopefully the NODE_LIST is ordered by NODE_RANK, this test is only expected to run with 2 nodes now - # We need to test with 4 nodes or more in the future, which should be easier with service discovery - gen_url = f"{get_the_other_host(0)}:8002" - with RemoteDisaggOpenAIServer(ctx_servers=[ctx_url], - gen_servers=[gen_url], - port=DISAGG_SERVER_PORT, - llmapi_launch=False, - env=env()) as server: - yield server - else: - print(f"skipping disagg_server for rank {RANK} node rank {NODE_RANK}") - yield None - - -@pytest.fixture(scope="module") -def client(disagg_server: RemoteDisaggOpenAIServer): - if is_pytest_node(): - return disagg_server.get_client() - else: - print(f"skipping client for rank {RANK} node rank {NODE_RANK}") - return None - - -def test_completion(client: openai.OpenAI, - disagg_server: RemoteDisaggOpenAIServer, model_name: str): - if len(NODE_LIST) != 2: - pytest.skip("This test is only expected to run with 2 nodes") - return - if is_pytest_node(): - print(f"running test_completion on rank {RANK} node rank {NODE_RANK}") - prompt = "What is the result of 1+1? Answer in one word: " - completion = client.completions.create( - model=model_name, - prompt=prompt, - max_tokens=10, - temperature=0.0, - ) - print(f"Output: {completion.choices[0].text}") - assert completion.id is not None - message = completion.choices[0].text - assert message.startswith('2.') - - # sleep 10 seconds to ensure a successful wait_for_endpoint_ready on rank1 - time.sleep(10) - disagg_server.terminate() - - elif is_gen_node(): - # keep gen workers alive until the test ends, again we hope the NODE_LIST is ordered by NODE_RANK - url = f"http://{get_the_other_host(0)}:{DISAGG_SERVER_PORT}/health/" - wait_for_endpoint_ready(url) - wait_for_endpoint_down(url) - assert True - else: - assert True diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py deleted file mode 100644 index 780844f2a77d..000000000000 --- a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py +++ /dev/null @@ -1,223 +0,0 @@ -import os -import shutil -import subprocess -import tempfile -import uuid - -import openai -import pytest - -from tensorrt_llm._utils import get_free_port -from tensorrt_llm.llmapi.disagg_utils import ServerRole - -from ..test_llm import get_model_path -from .openai_server import RemoteDisaggOpenAIServer, RemoteOpenAIServer -from .utils import expand_slurm_nodelist, wait_for_endpoint_down, wait_for_endpoint_ready - -RANK = int(os.environ.get("SLURM_PROCID", 0)) -NODE_RANK = int(os.environ.get("SLURM_NODEID", 0)) -NODE_LIST = expand_slurm_nodelist(os.environ.get("SLURM_NODELIST", "")) -SLURM_NTASKS_PER_NODE = int(os.environ.get("SLURM_NTASKS_PER_NODE", 1)) - -# This a multi-node QA test, use a fixed port instead of finding a free port -# so that all nodes can have the same disagg server config -DISAGG_SERVER_PORT = 8000 - - -# This test is supposed to run with 2 nodes or more -def is_ctx_node(): - assert len(NODE_LIST) == 2 - return NODE_RANK == 0 - - -def is_gen_node(): - assert len(NODE_LIST) == 2 - return NODE_RANK == 1 - - -def is_disagg_node(): - return NODE_RANK == 0 - - -# The test is run on multinodes but only the first node's output is used for assertion -def is_pytest_node(): - return NODE_RANK == 0 - - -def env(): - # Remove MPI related environment variables to isolate the ctx/gen processes - # so that they will not be in the same MPI communicator, otherwise the rank and world_size may mismatch - return { - k: v - for k, v in os.environ.items() - if not ("PMI_" in k or "OMPI_" in k or "PMIX_" in k or "SLURM_" in k) - and k not in ["UCX_TLS", "UCX_NET_DEVICES"] - } - - -@pytest.fixture -def model_name(): - return "llama-3.1-model/Llama-3.1-8B-Instruct" - - -@pytest.fixture -def disagg_host(): - return NODE_LIST[0] - - -@pytest.fixture(params=["etcd", "http"]) -def service_discovery(request, disagg_host: str): - if request.param == "etcd": - work_dir = tempfile.mkdtemp() - data_dir = f"{work_dir}/disagg_test-etcd-{uuid.uuid4()}" - etcd = subprocess.Popen( - [ - "etcd", - "--data-dir", - data_dir, - "--listen-client-urls", - "http://0.0.0.0:2379", - "--advertise-client-urls", - f"http://{disagg_host}:2379", - ] - ) - yield etcd, f"etcd://{disagg_host}:2379" - try: - etcd.kill() - etcd.wait(timeout=10) - shutil.rmtree(data_dir) - except Exception: - pass - else: - yield None, f"http://{disagg_host}:{DISAGG_SERVER_PORT}" - - -@pytest.fixture -def disagg_cluster_config(service_discovery: tuple): - _, uri = service_discovery - return { - "cluster_uri": uri, - "cluster_name": "", - } - - -@pytest.fixture -def worker(model_name: str, disagg_cluster_config: dict): - extra_config = { - "disagg_cluster": disagg_cluster_config, - "cache_transceiver_config": {"backend": "DEFAULT"}, - "kv_cache_config": { - "free_gpu_memory_fraction": 0.5, - "enable_block_reuse": False, - }, - "disable_overlap_scheduler": True, - } - # start workers on 0.0.0.0:, then the workers should be able to - # report their correct hostname:port to the disagg server - port = get_free_port() - if is_ctx_node(): - print(f"starting ctx_server for rank {RANK} node rank {NODE_RANK}") - model_path = get_model_path(model_name) - tp_size, pp_size = 1, 1 - args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] - with RemoteOpenAIServer( - model_path, - port=port, - cli_args=args, - host="0.0.0.0", - env=env(), - llmapi_launch=False, - rank=RANK % SLURM_NTASKS_PER_NODE, - extra_config=extra_config, - role=ServerRole.CONTEXT, - ) as server: - yield server - elif is_gen_node(): - print(f"starting gen_server for rank {RANK} node rank {NODE_RANK}") - model_path = get_model_path(model_name) - tp_size, pp_size = 1, 1 - args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] - with RemoteOpenAIServer( - model_path, - port=port, - cli_args=args, - host="0.0.0.0", - env=env(), - llmapi_launch=False, - rank=RANK % SLURM_NTASKS_PER_NODE, - extra_config=extra_config, - role=ServerRole.GENERATION, - ) as server: - yield server - else: - yield None - - -# different from non-service-discovery version, disagg server doesn't have to -# wait for ctx/gen servers to get ready -@pytest.fixture -def disagg_server(disagg_cluster_config: dict): - if is_disagg_node(): - disagg_config = { - "disagg_cluster": disagg_cluster_config, - "port": DISAGG_SERVER_PORT, - "hostname": "0.0.0.0", - } - print(f"starting disagg_server for rank {RANK} node rank {NODE_RANK}") - # ctx/gen servers are unnecessary for service discovery test - with RemoteDisaggOpenAIServer( - ctx_servers=[], - gen_servers=[], - port=DISAGG_SERVER_PORT, - disagg_config=disagg_config, - llmapi_launch=False, - env=env(), - wait_ready=False, # wait it to be ready in test body - ) as server: - yield server - else: - print(f"skipping disagg_server for rank {RANK} node rank {NODE_RANK}") - yield None - - -@pytest.fixture -def client(disagg_server: RemoteDisaggOpenAIServer): - if is_pytest_node(): - return disagg_server.get_client() - else: - print(f"skipping client for rank {RANK} node rank {NODE_RANK}") - return None - - -def test_completion( - disagg_server: RemoteDisaggOpenAIServer, - worker: RemoteOpenAIServer, - client: openai.OpenAI, - disagg_host: str, - model_name: str, -): - disagg_health_url = f"http://{disagg_host}:{DISAGG_SERVER_PORT}/health/" - wait_for_endpoint_ready(disagg_health_url) - if is_pytest_node(): - print(f"running test_completion on rank {RANK} node rank {NODE_RANK}") - prompt = "What is the result of 1+1? Answer in one word: " - for _ in range(10): - completion = client.completions.create( - model=model_name, - prompt=prompt, - max_tokens=10, - temperature=0.0, - ) - print(f"Output: {completion.choices[0].text}") - assert completion.id is not None - message = completion.choices[0].text - assert message.startswith("2.") - - disagg_server.terminate() - - elif is_gen_node(): - # keep gen workers alive until the test ends - wait_for_endpoint_down(disagg_health_url) - assert True - else: - assert True diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_guided_decoding.py b/tests/unittest/llmapi/apps/_test_openai_chat_guided_decoding.py index 22c04857bec0..ee8223f5a9a2 100644 --- a/tests/unittest/llmapi/apps/_test_openai_chat_guided_decoding.py +++ b/tests/unittest/llmapi/apps/_test_openai_chat_guided_decoding.py @@ -23,7 +23,6 @@ @pytest.fixture(scope="module", params=[ - "meta-llama/Llama-3.1-8B-Instruct", "openai/gpt-oss-120b", pytest.param("zai-org/GLM-5-FP8", marks=pytest.mark.skip_less_device(8)), @@ -58,9 +57,7 @@ def temp_extra_llm_api_options_file(model_name: str): @pytest.fixture(scope="module") def server(model_name: str, temp_extra_llm_api_options_file: str): - if model_name == "meta-llama/Llama-3.1-8B-Instruct": - model_path = get_model_path("llama-3.1-model/Llama-3.1-8B-Instruct") - elif model_name == "openai/gpt-oss-120b": + if model_name == "openai/gpt-oss-120b": model_path = get_model_path("gpt_oss/gpt-oss-120b") elif model_name == "zai-org/GLM-5-FP8": model_path = get_model_path("GLM-5-FP8") diff --git a/tests/unittest/llmapi/apps/_test_openai_multi_nodes.py b/tests/unittest/llmapi/apps/_test_openai_multi_nodes.py deleted file mode 100644 index 7413745e51a4..000000000000 --- a/tests/unittest/llmapi/apps/_test_openai_multi_nodes.py +++ /dev/null @@ -1,257 +0,0 @@ -import asyncio -import os -import re -import time - -import openai -import pytest -import torch -from utils.util import skip_num_gpus_less_than, skip_nvlink_inactive - -from ..test_llm import get_model_path, prompts -from .openai_server import RemoteOpenAIServer - -RANK = os.environ.get("SLURM_PROCID", 0) -MESSAGES = [{ - "role": "user", - "content": "Hello! How are you?" -}, { - "role": "assistant", - "content": "Hi! I am quite well, how can I help you today?" -}, { - "role": "user", - "content": "A song on old age?" -}] - - -@pytest.fixture(scope="module") -def model_name(): - return "llama-3.1-model/Llama-3.1-8B-Instruct" - - -@pytest.fixture(scope="module", params=['pytorch'], ids=["pytorch"]) -def backend(request): - return request.param - - -@pytest.fixture(scope="module", - params=[(16, 1), (8, 2)], - ids=lambda tp_pp_size: f'tp{tp_pp_size[0]}pp{tp_pp_size[1]}') -def tp_pp_size(request): - return request.param - - -@pytest.fixture(scope="module") -def server(model_name: str, backend: str, tp_pp_size: tuple): - os.environ["FORCE_DETERMINISTIC"] = "1" - model_path = get_model_path(model_name) - tp_size, pp_size = tp_pp_size - device_count = torch.cuda.device_count() - args = [ - "--tp_size", - f"{tp_size}", - "--pp_size", - f"{pp_size}", - "--gpus_per_node", - f"{device_count}", - "--kv_cache_free_gpu_memory_fraction", - "0.95", - "--backend", - backend, - ] - with RemoteOpenAIServer(model_path, args, llmapi_launch=True, - port=8001) as remote_server: - yield remote_server - - os.environ.pop("FORCE_DETERMINISTIC") - - -@pytest.fixture(scope="module") -def client(server: RemoteOpenAIServer): - return server.get_client() - - -@pytest.fixture(scope="module") -def async_client(server: RemoteOpenAIServer): - return server.get_async_client() - - -@skip_num_gpus_less_than(4) -def test_chat(client: openai.OpenAI, model_name: str): - if RANK == "0": - messages = [{ - "role": "system", - "content": "you are a helpful assistant" - }, { - "role": "user", - "content": "What is the result of 1+1? Answer in one word: " - }] - chat_completion = client.chat.completions.create( - model=model_name, - messages=messages, - max_tokens=1, - ) - assert chat_completion.id is not None - assert len(chat_completion.choices) == 1 - assert chat_completion.usage.completion_tokens == 1 - message = chat_completion.choices[0].message - - print(f"Output: {message.content}") - assert message.content == 'Two' - else: - time.sleep(30) - assert True - - -@skip_num_gpus_less_than(4) -def test_completion(client: openai.OpenAI, model_name: str): - if RANK == "0": - completion = client.completions.create( - model=model_name, - prompt=prompts, - max_tokens=5, - temperature=0.0, - ) - assert completion.choices[0].text == " D E F G H" - else: - time.sleep(30) - assert True - - -@skip_num_gpus_less_than(4) -@pytest.mark.asyncio(loop_scope="module") -async def test_chat_streaming(async_client: openai.AsyncOpenAI, - model_name: str): - if RANK == "0": - messages = [{ - "role": "system", - "content": "you are a helpful assistant" - }, { - "role": "user", - "content": "What is the result of 1+1? Answer in one word: " - }] - stream = await async_client.chat.completions.create( - model=model_name, - messages=messages, - max_tokens=1, - stream=True, - ) - async for chunk in stream: - delta = chunk.choices[0].delta - if delta.role: - assert delta.role == "assistant" - if delta.content: - assert delta.content == "Two" - else: - time.sleep(30) - assert True - - -@skip_num_gpus_less_than(4) -@pytest.mark.asyncio(loop_scope="module") -async def test_completion_streaming(async_client: openai.AsyncOpenAI, - model_name: str): - if RANK == "0": - completion = await async_client.completions.create( - model=model_name, - prompt=prompts, - max_tokens=5, - temperature=0.0, - stream=True, - ) - str_chunk = [] - async for chunk in completion: - str_chunk.append(chunk.choices[0].text) - assert "".join(str_chunk) == " D E F G H" - else: - time.sleep(30) - assert True - - -@skip_nvlink_inactive -@skip_num_gpus_less_than(4) -@pytest.mark.asyncio(loop_scope="module") -@pytest.mark.skip(reason="https://nvbugs/5112075") -async def test_multi_consistent_sync_chat(client: openai.OpenAI, - model_name: str): - """ - RCCA: https://nvbugs/4829393 - """ - if RANK == 0: - unique_content = set() - - async def send_request(messages=None): - try: - completion = client.chat.completions.create( - model=model_name, - messages=messages, - n=1, - max_tokens=1024, - temperature=0, - frequency_penalty=1.0, - stream=False, - stop=["hello"]) - unique_content.add(completion.choices[0].message.content) - except Exception as e: - print(f"Error: {e}") - - tasks = [] - for _ in range(50): - tasks.append(asyncio.create_task(send_request(MESSAGES))) - await asyncio.sleep(1) - - await asyncio.gather(*tasks) - - print(f"Number of unique responses: {len(unique_content)}") - assert len(unique_content) == 1, "Responses are not consistent" - content = list(unique_content)[0] - pattern = re.compile(r'[^a-zA-Z0-9\s\'\"]{5,}') - assert not bool(pattern.search(content)), content - else: - time.sleep(60) - assert True - - -@skip_nvlink_inactive -@skip_num_gpus_less_than(4) -@pytest.mark.asyncio(loop_scope="module") -@pytest.mark.skip(reason="https://nvbugs/5112075") -async def test_multi_consistent_async_chat(async_client: openai.AsyncOpenAI, - model_name: str): - """ - RCCA: https://nvbugs/4829393 - """ - - if RANK: - unique_content = set() - - async def send_request(messages=None): - try: - completion = await async_client.chat.completions.create( - model=model_name, - messages=messages, - n=1, - max_tokens=1024, - temperature=0, - frequency_penalty=1.0, - stream=False, - stop=["hello"]) - unique_content.add(completion.choices[0].message.content) - except Exception as e: - print(f"Error: {e}") - - tasks = [] - for _ in range(50): - tasks.append(asyncio.create_task(send_request(MESSAGES))) - await asyncio.sleep(1) - - await asyncio.gather(*tasks) - - print(f"Number of unique responses: {len(unique_content)}") - assert len(unique_content) == 1, "Responses are not consistent" - content = list(unique_content)[0] - pattern = re.compile(r'[^a-zA-Z0-9\s\'\"]{5,}') - assert not bool(pattern.search(content)), content - else: - time.sleep(60) - assert True diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_benchmark.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_benchmark.py index a9c1ecf2426a..7462711b0ef5 100644 --- a/tests/unittest/llmapi/apps/_test_trtllm_serve_benchmark.py +++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_benchmark.py @@ -58,11 +58,9 @@ def dataset_path(dataset_name: str): @skip_gpu_memory_less_than_80gb -@pytest.mark.parametrize("model_name", [ - "llama-3.1-model/Meta-Llama-3.1-8B", - pytest.param("gpt_oss/gpt-oss-20b", marks=skip_pre_hopper) -], - indirect=True) +@pytest.mark.parametrize( + "model_name", [pytest.param("gpt_oss/gpt-oss-20b", marks=skip_pre_hopper)], + indirect=True) def test_trtllm_serve_benchmark(server: RemoteOpenAIServer, benchmark_root: str, model_path: str): model_name = model_path.split("/")[-1] diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py index 84cf4a19c47f..66852ffa4677 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -251,7 +251,6 @@ def test_llm_with_kv_cache_retention_config(): (get_model_path('codellama/CodeLlama-7b-Instruct-hf'), False, 0.95), (llama_model_path, False, 0.95), (get_model_path(qwen3_tokenizer_model_name), False, 0.95), - (get_model_path('llama-3.1-model/Meta-Llama-3.1-8B'), False, 0.95), (get_model_path('DeepSeek-R1/DeepSeek-R1'), False, 0.95) ]) @pytest.mark.part0 diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index 0457ac706ef5..5a5d80c99382 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -35,9 +35,7 @@ sampling_params_for_aborting_request, run_llm_with_postprocess_parallel_and_result_handler, tinyllama_logits_processor_test_harness) -from utils.util import (force_ampere, similar, skip_fp8_pre_ada, - skip_gpu_memory_less_than_40gb, - skip_gpu_memory_less_than_80gb, +from utils.util import (force_ampere, similar, skip_gpu_memory_less_than_40gb, skip_gpu_memory_less_than_138gb, skip_ray) from utils.llm_data import llm_models_root from tensorrt_llm._torch.peft.lora.config import LoraConfig @@ -420,35 +418,6 @@ def test_nemotron_nas_lora(cuda_graph_config) -> None: llm.shutdown() -@skip_gpu_memory_less_than_80gb -@pytest.mark.part0 -@test_lora_with_and_without_cuda_graph -def test_llama_3_1_8b_fp8_with_bf16_lora(cuda_graph_config) -> None: - skip_fp8_pre_ada(use_fp8=True) - model_dir = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct-FP8" - lora_dir = f"{llm_models_root()}/lora/llama-3-chinese-8b-instruct-v2-lora" - prompt = "ē¾Žå›½ēš„首都ę˜Æå“Ŗ里ļ¼Ÿ" - reference = "华ē››é”æē‰¹åŒŗć€‚åŽē››é”æē‰¹åŒŗę˜Æē¾Žå›½ēš„首都和äø€äøŖč”Œę”æåŒŗ" - - lora_config = LoraConfig(lora_dir=[lora_dir], - max_lora_rank=64, - max_loras=2, - max_cpu_loras=2) - lora_req = LoRARequest("lora-chinese", 0, lora_dir) - - llm = LLM(model_dir, - lora_config=lora_config, - cuda_graph_config=cuda_graph_config) - - try: - output = llm.generate(prompt, - SamplingParams(max_tokens=20), - lora_request=[lora_req]) - finally: - llm.shutdown() - assert similar(output.outputs[0].text, reference) - - @pytest.mark.part2 @test_lora_with_and_without_cuda_graph def test_gemma3_1b_instruct_multi_lora(cuda_graph_config) -> None: diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index 88ec33dd7f77..ee2c79c70fde 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -497,16 +497,16 @@ def test_torch_compile_nodeids_are_private(): from test_common.session_reuse_hooks import _is_private_nodeid assert _is_private_nodeid( - "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::" - "test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=True]" + "accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::" + "test_fp8_prequantized[torch_compile=True]" ) assert _is_private_nodeid( "accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::" "test_nvfp4_multi_gpus_piecewise_cuda_graph[baseline]" ) assert not _is_private_nodeid( - "accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::" - "test_fp8[fp8kv=False-attn_backend=TRTLLM-torch_compile=False]" + "accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::" + "test_fp8_prequantized[torch_compile=False]" ) assert not _is_private_nodeid( "unittest/llmapi/test_llm_args.py::test_torch_compile_config_round_trip" diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index 7a9558ba3864..98290e877897 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -217,8 +217,8 @@ class TestConfigInfoMetrics: def test_model_config_info(self, collector): model_config = { - "model": "meta-llama/Llama-3.1-8B-Instruct", - "served_model_name": "Llama-3.1-8B-Instruct", + "model": "Qwen/Qwen3-8B", + "served_model_name": "Qwen3-8B", "dtype": "float16", "quantization": "none", "max_model_len": "4096", diff --git a/tests/unittest/scripts/test_check_model_registry.py b/tests/unittest/scripts/test_check_model_registry.py index d6b9c743bd91..3659ba496402 100644 --- a/tests/unittest/scripts/test_check_model_registry.py +++ b/tests/unittest/scripts/test_check_model_registry.py @@ -38,9 +38,9 @@ def mod(): def test_validate_models_allows_same_name_with_different_config_id(mod): models = [ - {"name": "meta-llama/Llama-3.1-8B-Instruct", "yaml_extra": ["world_size_1.yaml"]}, + {"name": "Qwen/Qwen3-8B", "yaml_extra": ["world_size_1.yaml"]}, { - "name": "meta-llama/Llama-3.1-8B-Instruct", + "name": "Qwen/Qwen3-8B", "config_id": "fp8", "yaml_extra": ["world_size_1.yaml", "fp8.yaml"], }, From 5a94cd7dcdd35dbc2c4086ba43c09438ec44b181 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:36:47 +0800 Subject: [PATCH 02/16] [None][fix] Remove leftover dead imports and stale test-list entries Cleanup fallout from the Llama-3.1-8B test removal conflict resolution: - Drop now-unused imports (llm_models_root, os, ExperimentConfig/main). - Add the sys/os imports actually needed by the sys.path.append call kept in test_sa.py. - Remove stale TestMiniMaxM2/TestMiniMaxM2_5 list entries: those classes were pruned upstream and were mistakenly reintroduced while merging a conflict block that legitimately needed to keep an unrelated, adjacent TestLlama4SpeculativeDecoding entry. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/test_lists/qa/llm_function_core.txt | 2 -- tests/unittest/_torch/speculative/hw_agnostic/test_sa.py | 3 ++- tests/unittest/_torch/speculative/test_eagle3.py | 1 - tests/unittest/auto_deploy/singlegpu/models/test_eagle.py | 1 - 4 files changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 6ba78745f57e..8dd23f691909 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -484,8 +484,6 @@ accuracy/test_llm_api_pytorch.py::TestLagunaXS_2_1::test_nvfp4_dflash accuracy/test_llm_api_pytorch.py::TestLlama4SpeculativeDecoding::test_guided_decoding TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestLlama4SpeculativeDecoding::test_llama4_eagle3[dynamic] TIMEOUT (60) accuracy/test_llm_api_pytorch.py::TestLlama4SpeculativeDecoding::test_llama4_eagle3[linear] TIMEOUT (60) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] -accuracy/test_llm_api_pytorch.py::TestMiniMaxM2_5::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index 5f4cc74e1343..7aa4bf23f847 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -1,8 +1,9 @@ +import os +import sys import unittest import pytest import torch -from utils.llm_data import llm_models_root from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.speculative.suffix_automaton import SAConfig, SuffixAutomatonManager diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 007a66428770..26379c1b8652 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -14,7 +14,6 @@ # limitations under the License. import json -import os import tempfile import unittest from contextlib import contextmanager diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py index 8bb0500803a0..d3f925185d97 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py @@ -33,7 +33,6 @@ ) __extra_import_path__ = ["~/examples/auto_deploy"] -from build_and_run_ad import ExperimentConfig, main from test_common.llm_data import hf_id_to_local_model_dir EAGLE_MODEL_HUB_ID = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" From 173924605385a638c33d59851f883833eb1f27da Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Fri, 11 Sep 2026 19:43:28 +0800 Subject: [PATCH 03/16] [None][fix] Fix waives.txt ordering Sort the MSA sparse-attention waive entries alphabetically to match the repo's required ordering. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 56f8d1e5097e..e6dc7b0ed37e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -261,10 +261,10 @@ test_e2e.py::test_ptp_quickstart_advanced[Nemotron-Nano-9B-v2-nvfp4-NVIDIA-Nemot test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (https://nvbugs/6605819) test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (bug pending, tracked in PR 17414) unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) -unittest/_torch/attention/sparse/msa/test_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype0] SKIP (https://nvbugs/6669902) unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype1] SKIP (https://nvbugs/6669902) unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py::test_sparse_decode_matches_msa_kernel SKIP (https://nvbugs/6669902) +unittest/_torch/attention/sparse/msa/test_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[exaone_moe_gqa_swa128-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6668773) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6641268) unittest/_torch/modeling/test_gemma4_e2e_dummy.py::test_e2e_text_31b_dummy SKIP (https://nvbugs/6607482) From e5d5c6a0b372e5fd91537226e9b619d239b93b8b Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Fri, 11 Sep 2026 21:41:53 +0800 Subject: [PATCH 04/16] [None][fix] Restore imports needed by the surviving Nemotron MTP smoke test test_super_mtp_smoke still needs hf_id_to_local_model_dir, ExperimentConfig, and main; these were dropped by mistake while resolving the Llama-3.1-8B removal merge conflict in this file. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- .../singlegpu/smoke/test_ad_speculative_decoding.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 3d287e3db6df..2c5d3ffe5a00 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -18,6 +18,10 @@ from tensorrt_llm.llmapi import MTPDecodingConfig +__extra_import_path__ = ["~/examples/auto_deploy"] +from build_and_run_ad import ExperimentConfig, main +from test_common.llm_data import hf_id_to_local_model_dir + def test_super_mtp_smoke(): """Test one-model MTP/Eagle runtime with a tiny Nemotron SuperV3 target.""" From 2061f0620533f0debf5bdaf90876391212e204c9 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sat, 12 Sep 2026 00:42:49 +0800 Subject: [PATCH 05/16] [None][revert] Drop tests/unittest/_torch and tests/unittest/auto_deploy Llama-3.1-8B changes Restore all files under these two directories to upstream/main, undoing the speculative-decoding and AutoDeploy smoke-test deletions from the Llama-3.1-8B cleanup. Narrows this branch's scope back to tests/integration/defs and the remaining tests/unittest files already touched elsewhere. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- .../sparse/rocketkv/test_rocketkv.py | 88 ++- .../_torch/modeling/test_modeling_llama.py | 714 ++++++++++++++++++ .../hw_agnostic/test_kv_cache_reuse.py | 84 +++ .../speculative/hw_agnostic/test_ngram.py | 68 ++ .../speculative/hw_agnostic/test_pard.py | 115 +++ .../_torch/speculative/hw_agnostic/test_sa.py | 208 ++++- .../speculative/hw_agnostic/test_spec_gate.py | 135 ++++ .../hw_agnostic/test_user_provided.py | 90 +++ .../_torch/speculative/test_eagle3.py | 609 ++++++++++++++- .../_utils_test/_model_test_utils.py | 11 +- .../smoke/test_ad_allreduce_strategies.py | 307 ++++++++ .../smoke/test_ad_build_small_multi.py | 49 ++ .../singlegpu/models/test_eagle.py | 145 +++- .../smoke/test_ad_build_small_single.py | 26 + .../smoke/test_ad_guided_decoding_regex.py | 73 ++ .../smoke/test_ad_speculative_decoding.py | 156 +++- .../singlegpu/smoke/test_ad_trtllm_serve.py | 117 +++ .../singlegpu/smoke/test_disagg.py | 59 ++ 18 files changed, 3044 insertions(+), 10 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_pard.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py create mode 100644 tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py create mode 100644 tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py create mode 100644 tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py create mode 100644 tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py diff --git a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py index 106a30b0ac14..4b93cd54dcea 100644 --- a/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py +++ b/tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py @@ -1,11 +1,14 @@ +import json import math +import os import pytest import torch +from utils.llm_data import llm_models_root from utils.util import getSMVersion import tensorrt_llm -from tensorrt_llm import SamplingParams +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs from tensorrt_llm._torch.attention.backends.sparse.rocket import ( RocketKVCacheManager, @@ -18,10 +21,86 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes from tensorrt_llm.bindings import DataType -from tensorrt_llm.llmapi import KvCacheConfig, RocketSparseAttentionConfig +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, RocketSparseAttentionConfig from tensorrt_llm.mapping import Mapping +@pytest.mark.skipif(getSMVersion() < 100, reason="RocketKV requires SM100 (Blackwell)") +@pytest.mark.parametrize("backend", ["pytorch"]) +@pytest.mark.parametrize("model_name", ["llama-3.1-model/Llama-3.1-8B-Instruct"]) +@pytest.mark.parametrize("attention_backend", ["VANILLA", "TRTLLM"]) +def test_model(backend, model_name, attention_backend, monkeypatch): + # RocketKV is a single-GPU path. Keep this test independent of MPI + # dynamic-process bootstrap so a cluster launch failure cannot mask the + # attention result. + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + model_dir = str(llm_models_root() / model_name) + max_batch_size = 16 + max_output_tokens = 128 + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.7, enable_block_reuse=False) + + kt_cache_dtype = "float8_e5m2" if attention_backend == "TRTLLM" else "bfloat16" + + sparse_attention_config = RocketSparseAttentionConfig( + window_size=32, + kernel_size=63, + prompt_budget=2048, + kt_cache_dtype=kt_cache_dtype, + ) + + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1, 2, 4, 8, 16], + enable_padding=True, + ) + + llm = LLM( + model=model_dir, + backend=backend, + kv_cache_config=kv_cache_config, + attn_backend=attention_backend, + sparse_attention_config=sparse_attention_config, + max_batch_size=max_batch_size, + max_seq_len=20480, + max_num_tokens=81920, + cuda_graph_config=None if attention_backend == "VANILLA" else cuda_graph_config, + ) + + inputs, references = [], [] + current_file = os.path.abspath(__file__) + current_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(current_file)))) + input_file = f"{current_dir}/multi_gpu/NIAH_simple_data.jsonl" + with open(input_file, "r") as f: + for line in f: + sample = json.loads(line) + inputs.append( + { + "prompt": sample["input_context"] + sample["input_query"], + } + ) + references.append(sample["outputs"][0]) + + with llm: + outputs = llm.generate( + inputs, + use_tqdm=True, + sampling_params=SamplingParams( + add_special_tokens=False, max_tokens=max_output_tokens, temperature=0.8, top_p=0.95 + ), + ) + + count = 0 + for ref, ret in zip(references, outputs): + print(f"ret: {ret.outputs[0].text}") + print(f"ref: {ref}") + if ref not in ret.outputs[0].text: + print(f"reference {ref} is not in the output {ret.outputs[0].text}") + else: + count = count + 1 + acc = count / len(outputs) + + assert acc >= 0.9, "accuracy test of rocketkv sparse attention failed" + + def create_rocket_kv_cache_manager( num_layers, num_kv_heads, @@ -693,6 +772,11 @@ def test_rocket_add_dummy_requests_forwards_capture_sampling_params(mocker): if __name__ == "__main__": + # RocketKV e2e tests + print("=== Testing RocketKV E2E tests ===") + test_model("pytorch", "llama-3.1-model/Llama-3.1-8B-Instruct", "VANILLA") + test_model("pytorch", "llama-3.1-model/Llama-3.1-8B-Instruct", "TRTLLM") + # Unit tests for sparse_kv_predict print("\n=== Testing sparse_kv_predict ===") test_sparse_kv_predict(1, 1) # bs=1, context only diff --git a/tests/unittest/_torch/modeling/test_modeling_llama.py b/tests/unittest/_torch/modeling/test_modeling_llama.py index c02e04053f20..c3fa5f427a5f 100644 --- a/tests/unittest/_torch/modeling/test_modeling_llama.py +++ b/tests/unittest/_torch/modeling/test_modeling_llama.py @@ -1,7 +1,66 @@ +import unittest +from copy import deepcopy from dataclasses import dataclass from typing import Any import torch +from _torch.helpers import create_mock_cuda_graph_runner +from parameterized import parameterized +from transformers import LlamaConfig +from transformers import LlamaForCausalLM as HFLlamaForCausalLM +from utils.llm_data import llm_models_root +from utils.util import default_dtype, getSMVersion + +import tensorrt_llm +from tensorrt_llm._torch.attention.backends.utils import get_attention_backend +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_llama import LlamaForCausalLM +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ + _update_kv_cache_draft_token_location +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.interface import (SpecMetadata, + SpeculativeDecodingMode) +from tensorrt_llm._torch.speculative.spec_tree_manager import SpecTreeManager +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + +LLAMA_3_1_8B_CONFIG = { + "architectures": ["LlamaForCausalLM"], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": 128001, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 131072, + "mlp_bias": False, + "model_type": "llama", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "rms_norm_eps": 1e-05, + "rope_scaling": { + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3" + }, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.43.0.dev0", + "use_cache": True, + "vocab_size": 128256 +} @dataclass(repr=False) @@ -23,3 +82,658 @@ def reduce_llama_config(mem_for_full_model: int, num_layers = int(config_dict["num_hidden_layers"] * model_fraction) num_layers = min(num_layers, default_num_layers) config_dict["num_hidden_layers"] = num_layers + + +class TestLlama(unittest.TestCase): + + @parameterized.expand([None, "FP8"]) + def test_llama_sanity(self, quant_algo): + config_dict = deepcopy(LLAMA_3_1_8B_CONFIG) + # 8B * sizeof(float16) plus some extra for activations + mem_for_full_model = (2 + 1) * 8 * 2**(30) + reduce_llama_config(mem_for_full_model, config_dict) + if config_dict["num_hidden_layers"] <= 0: + self.skipTest("Insufficient memory for a single Llama layer") + llama_config = LlamaConfig.from_dict(config_dict) + if quant_algo: + quant_config = QuantConfig(quant_algo=quant_algo) + else: + quant_config = None + if quant_algo == "FP8" and getSMVersion() < 89: + self.skipTest("This test is not supported in pre-Ada architecture") + + dtype = llama_config.torch_dtype + device = torch.device('cuda') + + with torch.device(device), default_dtype(dtype): + model_config = ModelConfig(pretrained_config=llama_config, + quant_config=quant_config) + llama = LlamaForCausalLM(model_config).to(device) + + input_ids = torch.tensor([100, 200, 300, 100, 200, 100, 400, 500], + dtype=torch.int, + device=device) + + context_sequence_lengths = [3, 2, 1] + sequence_lengths = context_sequence_lengths + [1, 1] + past_seen_tokens = [0, 0, 0, 62, 75] + request_ids = list(range(len(sequence_lengths))) + token_nums = (torch.tensor(past_seen_tokens) + + torch.tensor(sequence_lengths)).tolist() + prompt_lens = token_nums[:3] + past_seen_tokens[3:] + + num_blocks = 100 + tokens_per_block = 128 + head_dim = llama.config.hidden_size // llama.config.num_attention_heads + num_layers = llama.config.num_hidden_layers + num_kv_heads = llama.config.num_key_value_heads + max_seq_len = num_blocks * tokens_per_block + batch_size = len(context_sequence_lengths) + 2 + + if dtype == torch.half: + kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF + elif dtype == torch.bfloat16: + kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 + else: + raise ValueError("Invalid dtype") + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + kv_cache_config = KvCacheConfig(max_tokens=num_blocks * + tokens_per_block) + kv_cache_manager = KVCacheManager( + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + ) + kv_cache_manager.add_dummy_requests(request_ids, token_nums) + + metadata_cls = get_attention_backend(model_config.attn_backend).Metadata + attn_metadata = metadata_cls( + seq_lens=torch.tensor(sequence_lengths, dtype=torch.int), + num_contexts=len(context_sequence_lengths), + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=past_seen_tokens, + ), + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + max_num_requests=len(context_sequence_lengths) + 2, + max_num_tokens=8192, + ) + + position_ids = [] + for i, tokens in enumerate(past_seen_tokens): + seq_len = context_sequence_lengths[i] if i < len( + context_sequence_lengths) else 1 + position_id = torch.arange(tokens, + tokens + seq_len, + device=input_ids.device) + position_ids.append(position_id) + + position_ids = torch.cat(position_ids).unsqueeze(0) + + with torch.inference_mode(): + attn_metadata.prepare() + logits = llama.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata) + + self.assertEqual(len(past_seen_tokens), logits.shape[0]) + + with torch.inference_mode(): + attn_metadata.prepare() + logits = llama.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata, + return_context_logits=True) + self.assertEqual(input_ids.shape, logits.shape[:-1]) + + kv_cache_manager.shutdown() + + @parameterized.expand([ + Scenario(backend="VANILLA"), + Scenario(backend="FLASHINFER"), + Scenario(backend="FLASHINFER", use_cuda_graph=True), + Scenario(backend="TRTLLM"), + Scenario(backend="TRTLLM", use_cuda_graph=True), + ], lambda testcase_func, param_num, param: + f"{testcase_func.__name__}[{param.args[0]}]") + @torch.no_grad() + def test_llama_allclose_to_hf(self, scenario: Scenario) -> None: + """ + Compare output to HF + """ + backend = scenario.backend + metadata_cls = get_attention_backend(backend).Metadata + + torch.random.manual_seed(0) + config_dict = deepcopy(LLAMA_3_1_8B_CONFIG) + # 8B * sizeof(float16) plus some extra for activations + # times 2, since we'll need 2 of these + mem_for_full_model = (2 + 1) * 8 * 2**(30) * 4 + reduce_llama_config(mem_for_full_model, config_dict) + if config_dict["num_hidden_layers"] <= 0: + self.skipTest("Insufficient memory for a single Llama layer") + llama_config = LlamaConfig.from_dict(config_dict) + dtype = llama_config.torch_dtype + device = torch.device('cuda') + + with torch.device(device), default_dtype(dtype): + hf_llama = HFLlamaForCausalLM(llama_config).eval() + + model_config = ModelConfig(pretrained_config=llama_config, + attn_backend=backend) + + llama = LlamaForCausalLM(model_config).to(dtype).to(device) + llama.load_weights(hf_llama.state_dict()) + llama.post_load_weights() + + num_blocks = 1 + tokens_per_block = 128 + head_dim = llama.config.hidden_size // llama.config.num_attention_heads + num_layers = llama.config.num_hidden_layers + num_kv_heads = llama.config.num_key_value_heads + max_seq_len = num_blocks * tokens_per_block + batch_size = 1 + + if dtype == torch.half: + kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF + elif dtype == torch.bfloat16: + kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 + else: + raise ValueError("Invalid dtype") + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + kv_cache_config = KvCacheConfig(max_tokens=num_blocks * + tokens_per_block) + kv_cache_manager = KVCacheManager( + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + ) + + # context + input_ids = torch.tensor([100, 200, 300, 100, 200, 100, 400, 500], + dtype=torch.int, + device=device) + + num_cached_tokens_per_seq = [0] + request_ids = [1] + token_nums = [input_ids.size(-1)] + prompt_lens = [input_ids.size(-1)] + kv_cache_manager.add_dummy_requests(request_ids, token_nums) + + attn_metadata = metadata_cls( + seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + ) + + # Note: no CUDA graphs for prefill, the graph runner is built for + # decoding only. + position_ids = [torch.arange(0, input_ids.size(-1))] + position_ids = torch.cat(position_ids).unsqueeze(0).cuda() + with torch.inference_mode(): + attn_metadata.prepare() + logits = llama.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata) + ref = hf_llama.forward(input_ids=input_ids.unsqueeze(0), + position_ids=position_ids, + use_cache=True) + + torch.testing.assert_close(logits, + ref.logits[:, -1].float(), + atol=0.4, + rtol=0.4) + + # gen + gen_input_ids = torch.tensor([600], dtype=torch.int, device=device) + + num_cached_tokens_per_seq = [input_ids.size(-1)] + + attn_metadata = metadata_cls( + seq_lens=torch.tensor([gen_input_ids.size(-1)], dtype=torch.int), + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + ) + + gen_position_ids = [ + torch.arange(input_ids.size(-1), + input_ids.size(-1) + gen_input_ids.size(-1)) + ] + gen_position_ids = torch.cat(gen_position_ids).unsqueeze(0).cuda() + + graph_runner = create_mock_cuda_graph_runner( + 1) if scenario.use_cuda_graph else None + + def run_forward(input_ids, position_ids, attn_metadata): + attn_metadata.prepare() + if not scenario.use_cuda_graph: + return llama.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata) + else: + inputs = { + "input_ids": input_ids, + "position_ids": position_ids, + "attn_metadata": attn_metadata, + } + key = (1, 0, False) + graph_runner.capture(key, + lambda inputs: llama.forward(**inputs), + inputs) + for _ in range(2): + # Run it twice. This helps us catch problems if buffers are accidentally reallocated + # in prepare(). + attn_metadata.prepare() + logits = graph_runner.replay(key, inputs) + return logits + + if scenario.use_cuda_graph: + attn_metadata = attn_metadata.create_cuda_graph_metadata(1) + + with torch.inference_mode(): + logits = run_forward(input_ids=gen_input_ids, + position_ids=gen_position_ids, + attn_metadata=attn_metadata) + ref = hf_llama.forward(input_ids=gen_input_ids.unsqueeze(0), + position_ids=gen_position_ids, + past_key_values=ref.past_key_values, + use_cache=True) + + torch.testing.assert_close(logits, + ref.logits[:, -1].float(), + atol=0.4, + rtol=0.4) + if graph_runner is not None: + graph_runner.clear() + kv_cache_manager.shutdown() + + @torch.no_grad() + def test_llama_verification_with_kv_cache_relocation(self) -> None: + """ + Verify the output of the model with kv cache relocation + """ + backend = "TRTLLM" + metadata_cls = get_attention_backend(backend).Metadata + + config_dict = deepcopy(LLAMA_3_1_8B_CONFIG) + + llama_config = LlamaConfig.from_dict(config_dict) + dtype = llama_config.torch_dtype + device = torch.device('cuda') + + with torch.device(device), default_dtype(dtype): + models_path = llm_models_root() + model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + hf_llama = HFLlamaForCausalLM.from_pretrained( + model_dir, + torch_dtype=torch.bfloat16, + device_map="cuda", + ).eval() + + model_config = ModelConfig(pretrained_config=llama_config, + attn_backend=backend) + + llama = LlamaForCausalLM(model_config).to(dtype).to(device) + llama.load_weights(hf_llama.state_dict()) + + num_blocks = 2 + tokens_per_block = 32 + head_dim = llama.config.hidden_size // llama.config.num_attention_heads + num_layers = llama.config.num_hidden_layers + num_kv_heads = llama.config.num_key_value_heads + num_heads_per_kv = llama.config.num_attention_heads // num_kv_heads + max_seq_len = num_blocks * tokens_per_block + batch_size = 1 + + if dtype == torch.half: + kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF + elif dtype == torch.bfloat16: + kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 + else: + raise ValueError("Invalid dtype") + kv_cache_dtype_byte_size = 2 + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + kv_cache_config = KvCacheConfig(max_tokens=num_blocks * + tokens_per_block) + kv_cache_manager = KVCacheManager( + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + ) + + # context + input_ids = torch.tensor([ + 128000, 32, 6369, 1990, 264, 22999, 1217, 323, 459, 21075, 11478, + 18328, 13, 578, 18328, 6835, 11190, 11, 11944, 11, 323, 48887, + 11503, 311, 279, 1217, 596, 4860, 13, 14194, 25, 22691, 36660, 3931, + 2891, 25 + ], + dtype=torch.int, + device=device) + + num_cached_tokens_per_seq = [0] + request_ids = [900] + token_nums = [input_ids.size(-1)] + prompt_lens = [input_ids.size(-1)] + requests = kv_cache_manager.add_dummy_requests(request_ids, token_nums) + request = requests[0] + + attn_metadata = metadata_cls( + seq_lens=torch.tensor([input_ids.size(-1)], dtype=torch.int), + num_contexts=1, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + num_heads_per_kv=num_heads_per_kv, + ) + + position_ids = [torch.arange(0, input_ids.size(-1))] + position_ids = torch.cat(position_ids).unsqueeze(0).cuda() + with torch.inference_mode(): + attn_metadata.prepare() + logits = llama.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata) + + def run_forward(input_ids, position_ids, attn_metadata): + return llama.forward(input_ids=input_ids, + position_ids=position_ids, + attn_metadata=attn_metadata, + return_context_logits=True) + + # prepare for the first generation + gen_input_ids_0 = torch.tensor([ + 22691, 11, 0, 13, 15592, 323, 315, 12, 311, 362, 220, 32, 362, 426, + 330, 358, 362, 358, 358, 362, 32, 0, 13, 32, 6369 + ], + dtype=torch.int, + device=device) + spec_decoding_position_offsets = torch.tensor([ + 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 3 + ], + dtype=torch.int, + device=device) + spec_decoding_packed_mask = torch.tensor( + [ + 1, 3, 5, 9, 17, 33, 65, 129, 257, 513, 1025, 2051, 4099, 8195, + 16387, 32771, 65541, 131077, 262153, 524297, 1048593, 2097169, + 4194321, 8388641, 16842757 + ], + dtype=torch.int, + device=device).unsqueeze(0).unsqueeze(2) + + num_cached_tokens_per_seq = [input_ids.size(-1)] + is_spec_decoding_enabled = True + use_spec_decoding = True + is_spec_dec_tree = True + is_spec_dec_dynamic_tree = True + max_total_draft_tokens = gen_input_ids_0.size(-1) - 1 + + spec_tree_mgr = SpecTreeManager( + max_num_requests=1, + use_dynamic_tree=is_spec_dec_dynamic_tree, + max_total_draft_tokens=max_total_draft_tokens, + max_draft_len=max_total_draft_tokens, + eagle_choices=None, + dynamic_tree_max_topK=10, + ) + # Populate with test data + spec_tree_mgr.spec_dec_position_offsets[:1, :].copy_( + spec_decoding_position_offsets.unsqueeze(0)) + spec_tree_mgr.spec_dec_packed_mask[:1, :, :].copy_( + spec_decoding_packed_mask) + + attn_metadata_gen_phase_0 = metadata_cls( + seq_lens=torch.tensor([gen_input_ids_0.size(-1)], dtype=torch.int), + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + is_spec_decoding_enabled=is_spec_decoding_enabled, + use_spec_decoding=use_spec_decoding, + is_spec_dec_tree=is_spec_dec_tree, + is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, + num_heads_per_kv=num_heads_per_kv, + ) + attn_metadata_gen_phase_0.prepare() + attn_metadata_gen_phase_0.update_spec_dec_param( + batch_size=batch_size, + is_spec_decoding_enabled=is_spec_decoding_enabled, + is_spec_dec_dynamic_tree=is_spec_dec_dynamic_tree, + is_spec_dec_tree=is_spec_dec_tree, + max_draft_len=max_total_draft_tokens, + max_total_draft_tokens=max_total_draft_tokens, + spec_tree_manager=spec_tree_mgr, + ) + + gen_position_ids_0 = [ + torch.full((gen_input_ids_0.size(-1), ), + input_ids.size(-1), + dtype=torch.int64) + ] + gen_position_ids_0 = torch.cat(gen_position_ids_0).unsqueeze(0).cuda() + + with torch.inference_mode(): + gen_logits_0 = run_forward(input_ids=gen_input_ids_0, + position_ids=gen_position_ids_0, + attn_metadata=attn_metadata_gen_phase_0) + + request.py_num_accepted_draft_tokens = 1 + request.py_num_accepted_draft_tokens_indices = [1] + request.py_rewind_len = gen_input_ids_0.size( + -1) - request.py_num_accepted_draft_tokens - 1 + request.state = LlmRequestState.GENERATION_IN_PROGRESS + scheduled_requests = ScheduledRequests() + scheduled_requests.generation_requests = [request] + kv_cache_manager.max_draft_len = gen_input_ids_0.size(-1) - 1 + _update_kv_cache_draft_token_location(kv_cache_manager, + scheduled_requests, + attn_metadata_gen_phase_0, + kv_cache_dtype_byte_size) + if request.py_rewind_len > 0: + kv_cache_manager.rewind_kv_cache(request, request.py_rewind_len) + torch.cuda.synchronize() + + # prepare for the second generation + gen_input_ids_1 = torch.tensor([2650, 649], + dtype=torch.int, + device=device) + + num_cached_tokens_per_seq_1 = [ + input_ids.size(-1) + request.py_num_accepted_draft_tokens + 1 + ] + attn_metadata_gen_phase_0.seq_lens = torch.tensor( + [gen_input_ids_1.size(-1)], dtype=torch.int) + attn_metadata_gen_phase_0.kv_cache_params.num_cached_tokens_per_seq = num_cached_tokens_per_seq_1 + + attn_metadata_gen_phase_0.spec_decoding_position_offsets = None + attn_metadata_gen_phase_0.spec_decoding_packed_mask = None + attn_metadata_gen_phase_0.spec_decoding_generation_lengths = None + attn_metadata_gen_phase_0.prepare() + is_tree_phase1 = is_spec_dec_tree if get_sm_version() < 100 else False + spec_tree_mgr_phase1 = None + spec_metadata_phase1 = None + if is_tree_phase1: + max_draft_1 = gen_input_ids_1.size(-1) - 1 + spec_tree_mgr_phase1 = SpecTreeManager( + max_num_requests=1, + use_dynamic_tree=True, + max_total_draft_tokens=max_draft_1, + max_draft_len=max_draft_1, + eagle_choices=None, + dynamic_tree_max_topK=10, + ) + spec_metadata_phase1 = SpecMetadata( + max_num_requests=1, + max_draft_len=max_draft_1, + max_total_draft_tokens=max_draft_1, + spec_dec_mode=SpeculativeDecodingMode.EAGLE3, + ) + attn_metadata_gen_phase_0.update_spec_dec_param( + batch_size=batch_size, + is_spec_decoding_enabled=is_spec_decoding_enabled, + is_spec_dec_tree=is_tree_phase1, + is_spec_dec_dynamic_tree=is_tree_phase1, + max_draft_len=gen_input_ids_1.size(-1) - 1, + max_total_draft_tokens=gen_input_ids_1.size(-1) - 1, + spec_metadata=spec_metadata_phase1, + spec_tree_manager=spec_tree_mgr_phase1) + + gen_position_ids_1 = [ + torch.full( + (gen_input_ids_1.size(-1), ), + input_ids.size(-1) + request.py_num_accepted_draft_tokens + 1, + dtype=torch.int64) + ] + gen_position_ids_1 = torch.cat(gen_position_ids_1).unsqueeze(0).cuda() + + with torch.inference_mode(): + gen_logits_1 = run_forward(input_ids=gen_input_ids_1, + position_ids=gen_position_ids_1, + attn_metadata=attn_metadata_gen_phase_0) + + torch.cuda.synchronize() + + # prepare for the reference generation + gen_input_ids_ref = torch.tensor([22691, 0, 2650, 649], + dtype=torch.int, + device=device) + num_cached_tokens_per_seq_ref = [input_ids.size(-1)] + + attn_metadata_ref = metadata_cls( + seq_lens=torch.tensor([gen_input_ids_ref.size(-1)], + dtype=torch.int), + num_contexts=0, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=num_cached_tokens_per_seq_ref, + ), + max_num_requests=1, + max_num_tokens=8192, + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + is_spec_decoding_enabled=is_spec_decoding_enabled, + use_spec_decoding=use_spec_decoding, + is_spec_dec_tree=is_spec_dec_tree, + is_spec_dec_dynamic_tree=False, + num_heads_per_kv=num_heads_per_kv, + ) + + attn_metadata_ref.spec_decoding_position_offsets = None + attn_metadata_ref.spec_decoding_packed_mask = None + attn_metadata_ref.spec_decoding_generation_lengths = None + attn_metadata_ref.prepare() + is_tree_ref = is_spec_dec_tree if get_sm_version() < 100 else False + spec_tree_mgr_ref = None + spec_metadata_ref = None + if is_tree_ref: + max_draft_ref = gen_input_ids_ref.size(-1) - 1 + spec_tree_mgr_ref = SpecTreeManager( + max_num_requests=1, + use_dynamic_tree=True, + max_total_draft_tokens=max_draft_ref, + max_draft_len=max_draft_ref, + eagle_choices=None, + dynamic_tree_max_topK=10, + ) + spec_metadata_ref = SpecMetadata( + max_num_requests=1, + max_draft_len=max_draft_ref, + max_total_draft_tokens=max_draft_ref, + spec_dec_mode=SpeculativeDecodingMode.EAGLE3, + ) + attn_metadata_ref.update_spec_dec_param( + batch_size=batch_size, + is_spec_decoding_enabled=is_spec_decoding_enabled, + is_spec_dec_tree=is_tree_ref, + is_spec_dec_dynamic_tree=is_tree_ref, + max_draft_len=gen_input_ids_ref.size(-1) - 1, + max_total_draft_tokens=gen_input_ids_ref.size(-1) - 1, + spec_metadata=spec_metadata_ref, + spec_tree_manager=spec_tree_mgr_ref) + + gen_position_ids_ref = [ + torch.full((gen_input_ids_ref.size(-1), ), + input_ids.size(-1), + dtype=torch.int64) + ] + gen_position_ids_ref = torch.cat(gen_position_ids_ref).unsqueeze( + 0).cuda() + + with torch.inference_mode(): + gen_logits_ref = run_forward(input_ids=gen_input_ids_ref, + position_ids=gen_position_ids_ref, + attn_metadata=attn_metadata_ref) + + torch.cuda.synchronize() + torch.testing.assert_close(gen_logits_1[0, :], + gen_logits_ref[2, :], + atol=1.0, + rtol=1.0) + torch.testing.assert_close(gen_logits_1[1, :], + gen_logits_ref[3, :], + atol=1.0, + rtol=1.0) + + token_id_ref = torch.argmax(gen_logits_ref[3, :], dim=-1) + token_id_gen = torch.argmax(gen_logits_1[1, :], dim=-1) + assert token_id_ref == token_id_gen, "Greedy sampling token id not match" + + kv_cache_manager.shutdown() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py new file mode 100644 index 000000000000..1ac2e3a5948a --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kv_cache_reuse.py @@ -0,0 +1,84 @@ +# 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 unittest + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig + + +@pytest.mark.high_cuda_memory +def test_eagle3_one_model_kv_cache_reuse() -> None: + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + enable_partial_reuse=True, + use_kv_cache_manager_v2=True, + tokens_per_block=32, + max_tokens=8192, + ) + llm_common_config = dict( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1]), + max_batch_size=1, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + ) + llm_spec = LLM( + **llm_common_config, + speculative_config=Eagle3DecodingConfig( + max_draft_len=4, + speculative_model=eagle_model_dir, + eagle3_one_model=True, + ), + ) + prompt = ("The quick brown fox jumped over the lazy dog. " * 20).strip() + " Once upon a time," + poison_prompt = "X Y Z W A B C D E F G H I J K L M N O P Q R S T U V " * 60 + sampling_params = SamplingParams( + max_tokens=64, + temperature=0, + ignore_eos=True, + return_perf_metrics=True, + ) + + llm_spec.generate(poison_prompt, SamplingParams(max_tokens=8, temperature=0)) + cold_result = llm_spec.generate(prompt, sampling_params) + reuse_result = llm_spec.generate(prompt, sampling_params) + llm_spec.shutdown() + + cold_acceptance = float(cold_result.avg_decoded_tokens_per_iter) + reuse_acceptance = float(reuse_result.avg_decoded_tokens_per_iter) + assert cold_result.outputs[0].text == reuse_result.outputs[0].text + assert cold_acceptance > 1.0 + assert reuse_acceptance > 1.0 + assert reuse_result.cached_tokens >= 64 + assert reuse_acceptance >= cold_acceptance - 0.2 + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py new file mode 100644 index 000000000000..adc87b4f23f8 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ngram.py @@ -0,0 +1,68 @@ +import unittest + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, NGramDecodingConfig + + +# TODO: add disable_overlap_scheduler=False +@pytest.mark.parametrize( + "disable_overlap_scheduler,use_cuda_graph,attn_backend", + [[True, False, "TRTLLM"], [True, True, "TRTLLM"], [True, False, "FLASHINFER"]], +) +@pytest.mark.high_cuda_memory +def test_llama_ngram(disable_overlap_scheduler: bool, use_cuda_graph: bool, attn_backend: str): + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 20: + pytest.skip("Not enough memory to load target model") + + max_batch_size = 2 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) + cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=llm_models_root() / "llama-3.1-model" / "Meta-Llama-3.1-8B", + backend="pytorch", + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_num_tokens=2048, + ) + + spec_config = NGramDecodingConfig( + max_draft_len=max_draft_len, + max_matching_ngram_size=2, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=True, + ) + + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) + + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(**llm_common_config) + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): + # The spec decode algorithm currently guarantees identical results + assert text_spec == text_ref + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py new file mode 100644 index 000000000000..5e93a70b174c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_pard.py @@ -0,0 +1,115 @@ +import unittest + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, PARDDecodingConfig + + +@pytest.mark.parametrize("disable_overlap_scheduler", [True, False]) +def test_pard(disable_overlap_scheduler: bool): + """Test PARD speculative decoding with CUDA graph support. + + This test verifies that PARD (Parallel Draft) speculative decoding works + correctly with CUDA graphs and padding enabled. + """ + attn_backend = "TRTLLM" + enable_block_reuse = False + enable_chunked_prefill = False + + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + # Test with 3 requests and max_batch_size=4 to trigger padding + max_batch_size = 4 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, max_tokens=2048) + use_cuda_graph = True + cuda_graph_config = ( + CudaGraphConfig(batch_sizes=[1, 2, 4], enable_padding=True) if use_cuda_graph else None + ) + + llm_common_config = dict( + model=target_model_dir, + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=2048, + enable_chunked_prefill=enable_chunked_prefill, + ) + + spec_config = PARDDecodingConfig( + max_draft_len=max_draft_len, + speculative_model=pard_model_dir, + ) + + # Create the LLM instance + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + prompts = [ + "The capital of France is", + "The president of the United States is", + "The future of AI is", + ] + + sampling_params = SamplingParams(max_tokens=1024, temperature=0) + llm_spec.generate(prompts, sampling_params) + llm_spec.shutdown() + + +@pytest.mark.high_cuda_memory +def test_pard_rejection(): + """PARD with rejection sampling on: the block-capture rejection path + (draft-prob scatter -> fail-closed guard -> rejection acceptance) runs + end-to-end with non-greedy sampling and produces coherent output.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + pard_model_dir = f"{models_path}/PARD-Llama-3.2-1B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + spec_config = PARDDecodingConfig( + max_draft_len=4, + speculative_model=pard_model_dir, + use_rejection_sampling=True, + ) + llm_spec = LLM( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2], enable_padding=True), + max_batch_size=2, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=2048), + max_seq_len=2048, + speculative_config=spec_config, + ) + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + # Non-greedy so rejection sampling actually engages (all-greedy bypasses it). + sampling_params = SamplingParams( + max_tokens=32, temperature=0.8, top_p=0.95, top_k=50, seed=1234 + ) + outputs = llm_spec.generate(prompts, sampling_params) + llm_spec.shutdown() + + assert len(outputs) == len(prompts) + for out in outputs: + assert len(out.outputs[0].token_ids) > 0 + assert out.outputs[0].text.strip() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index 7aa4bf23f847..979941d79e86 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -1,15 +1,215 @@ -import os -import sys import unittest import pytest import torch +from utils.llm_data import llm_models_root +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._torch.speculative.suffix_automaton import SAConfig, SuffixAutomatonManager -from tensorrt_llm.llmapi import SADecodingConfig +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SADecodingConfig + + +def get_perf_metrics(result): + """Extract performance metrics from result using built-in request_perf_metrics.""" + metrics = {} + if result.outputs and result.outputs[0].request_perf_metrics: + perf = result.outputs[0].request_perf_metrics + timing = perf.timing_metrics + # Convert timedelta to seconds + metrics["arrival_time"] = timing.arrival_time.total_seconds() + metrics["first_token_time"] = timing.first_token_time.total_seconds() + metrics["last_token_time"] = timing.last_token_time.total_seconds() + # Calculate TTFT and E2E latency + metrics["ttft"] = metrics["first_token_time"] - metrics["arrival_time"] + metrics["e2e"] = metrics["last_token_time"] - metrics["arrival_time"] + return metrics + + +# Test parameter combinations: +# - disable_overlap_scheduler: Controls scheduler mode (False=overlap enabled) +# - use_cuda_graph: Whether to use CUDA graph capture +# - attn_backend: Attention implementation (TRTLLM only - FLASHINFER not supported) +# - max_matching_ngram_size: SA matching mode (2=fixed size, -1=longest match) +# +# NOTE: FLASHINFER target decode supports multiple queries per request, but +# non-shared one-engine modes still require a separate draft KV cache. The +# draft KV metadata/manager swap is currently implemented only for TRTLLM +# attention. Shared-target-KV modes use a separate FlashInfer metadata view. +@pytest.mark.parametrize( + "disable_overlap_scheduler,use_cuda_graph,attn_backend,max_matching_ngram_size", + [ + [False, False, "TRTLLM", 2], + [False, True, "TRTLLM", 2], + [True, False, "TRTLLM", 2], + [True, True, "TRTLLM", 2], + [False, False, "TRTLLM", -1], + ], +) +@pytest.mark.high_cuda_memory +def test_llama_sa( + disable_overlap_scheduler: bool, + use_cuda_graph: bool, + attn_backend: str, + max_matching_ngram_size: int, +): + """Test SA (Suffix Automaton) speculative decoding correctness and acceptance rate. + + Verifies: + 1. Speculative decoding produces identical results to baseline + 2. SA drafting produces draft tokens that get accepted + 3. Multi-token acceptance occurs (acceptanceLength > 1) + """ + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 20: + pytest.skip("Not enough memory to load target model") + + print( + f"\nTest config: disable_overlap_scheduler={disable_overlap_scheduler}, " + f"use_cuda_graph={use_cuda_graph}, attn_backend={attn_backend}, " + f"max_matching_ngram_size={max_matching_ngram_size}" + ) + + max_batch_size = 1 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) + cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=llm_models_root() / "llama-3.1-model" / "Meta-Llama-3.1-8B", + backend="pytorch", + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_num_tokens=2048, + enable_iter_perf_stats=True, + ) + + spec_config = SADecodingConfig( + max_draft_len=max_draft_len, + max_matching_ngram_size=max_matching_ngram_size, + ) + + # Use prompts that encourage repetitive patterns for better SA/ngram matching + prompts = [ + "Count from 1 to 50: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, " + "16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, " + "34, 35,", + ] + # Enable perf metrics collection via return_perf_metrics=True + sampling_params = SamplingParams( + max_tokens=64, ignore_eos=True, temperature=0, return_perf_metrics=True + ) + + # Run with speculative decoding + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + + # Get spec decoding stats before shutdown + stats = llm_spec.get_stats(timeout=5) + iterations_with_spec = [] + for stat in stats: + if "specDecodingStats" in stat: + spec_stats = stat["specDecodingStats"] + if spec_stats.get("numDraftTokens", 0) > 0: + iterations_with_spec.append(spec_stats) + + # Get perf metrics using built-in request_perf_metrics + spec_metrics = get_perf_metrics(results_spec[0]) if results_spec else {} + + llm_spec.shutdown() + + # Run reference without speculative decoding + llm_ref = LLM(**llm_common_config) + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] -sys.path.append(os.path.join(os.path.dirname(__file__), "..")) + # Get perf metrics for reference + ref_metrics = get_perf_metrics(results_ref[0]) if results_ref else {} + + llm_ref.shutdown() + + # Verify 1: Identical results (correctness) + for i, (text_spec, text_ref) in enumerate(zip(generated_text_spec, generated_text_ref)): + assert text_spec == text_ref, ( + f"Prompt {i}: Spec decode result differs from baseline.\n" + f"Spec: {text_spec}\nRef: {text_ref}" + ) + print("Correctness verified: spec decode matches baseline") + + # Verify 2: Spec decoding stats show drafting occurred + assert len(iterations_with_spec) > 0, ( + f"SA should have iterations with specDecodingStats. " + f"Got {len(stats)} total stats but 0 with draft tokens." + ) + + total_draft = sum(s["numDraftTokens"] for s in iterations_with_spec) + total_accepted = sum(s["numAcceptedTokens"] for s in iterations_with_spec) + avg_acceptance_len = sum(s["acceptanceLength"] for s in iterations_with_spec) / len( + iterations_with_spec + ) + + print("Spec decoding stats:") + print(f" Iterations with drafting: {len(iterations_with_spec)}") + print(f" Total draft tokens: {total_draft}") + print(f" Total accepted tokens: {total_accepted}") + print(f" Average acceptance length: {avg_acceptance_len:.2f}") + print(f" Acceptance rate: {total_accepted / total_draft * 100:.1f}%") + + assert total_draft > 0, "SA should produce draft tokens" + assert total_accepted > 0, ( + f"SA should accept some draft tokens. " + f"Got {total_accepted} accepted out of {total_draft} drafted" + ) + + # Verify 3: Multi-token acceptance (acceptanceLength > 1) + has_multi_token_acceptance = any(s["acceptanceLength"] > 1.0 for s in iterations_with_spec) + print(f" Has multi-token acceptance: {has_multi_token_acceptance}") + + assert has_multi_token_acceptance, ( + "Expected at least one iteration with acceptanceLength > 1 for repetitive pattern" + ) + + # Print performance comparison using built-in metrics + print("\n" + "=" * 70) + print("PERFORMANCE COMPARISON (using request_perf_metrics)") + print("=" * 70) + print( + f"Config: overlap_scheduler={'enabled' if not disable_overlap_scheduler else 'disabled'}, " + f"cuda_graph={'enabled' if use_cuda_graph else 'disabled'}" + ) + print("-" * 70) + print(f"{'Metric':<30} {'Spec Decoding':<20} {'Reference':<20}") + print("-" * 70) + + # Print TTFT (Time to First Token) + ttft_spec = spec_metrics.get("ttft", None) + ttft_ref = ref_metrics.get("ttft", None) + ttft_spec_str = f"{ttft_spec * 1000:.2f} ms" if ttft_spec else "N/A" + ttft_ref_str = f"{ttft_ref * 1000:.2f} ms" if ttft_ref else "N/A" + print(f"{'TTFT':<30} {ttft_spec_str:<20} {ttft_ref_str:<20}") + + # Print E2E latency + e2e_spec = spec_metrics.get("e2e", None) + e2e_ref = ref_metrics.get("e2e", None) + e2e_spec_str = f"{e2e_spec * 1000:.2f} ms" if e2e_spec else "N/A" + e2e_ref_str = f"{e2e_ref * 1000:.2f} ms" if e2e_ref else "N/A" + print(f"{'E2E Latency':<30} {e2e_spec_str:<20} {e2e_ref_str:<20}") + + # Calculate and print speedup + if e2e_spec and e2e_ref and e2e_spec > 0: + speedup = e2e_ref / e2e_spec + print("-" * 70) + print(f"{'Speedup (E2E)':<30} {speedup:.2f}x") + print("=" * 70 + "\n") + + # Synchronize CUDA to catch any async memory errors before test completes. + # This ensures errors are attributed to this test rather than propagating + # to subsequent tests. + torch.cuda.synchronize() @pytest.mark.cpu_only diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py index e1486a8d9696..b523ed443727 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_gate.py @@ -1,8 +1,143 @@ import unittest +from unittest.mock import patch import pytest +import torch +from utils.llm_data import llm_models_root +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.speculative.speculation_gate import SpeculationGate +from tensorrt_llm.llmapi import CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig +from tensorrt_llm.logger import logger + + +@pytest.fixture(scope="function") +def enforce_single_worker(monkeypatch): + """Mock functions don't work with multiple processes, so we enforce single worker.""" + monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") + yield + + +# Tests that the SpeculationGate correctly disables speculative decoding +# when the average acceptance rate drops below the threshold. +# This test uses a mock to simulate low acceptance rates and verifies +# that the spec gate triggers and disables speculation. +@pytest.mark.high_cuda_memory +def test_spec_gate_e2e(enforce_single_worker): + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + models_path = llm_models_root() + eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + max_batch_size = 3 + max_draft_len = 4 + acceptance_rate_window_size = 3 + acceptance_rate_threshold = 0.6 + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=0.6, + ) + cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) + + llm_common_config = dict( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=False, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + cuda_graph_config=cuda_graph_config, + enable_chunked_prefill=False, + max_num_tokens=8192, + ) + + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=eagle_model_dir, + eagle3_one_model=True, + acceptance_rate_window_size=acceptance_rate_window_size, + acceptance_rate_threshold=acceptance_rate_threshold, + ) + + prompts = [ + "The capital of France is", + "The president of the United States is", + "What is the capital of Australia?", + ] + sampling_params = SamplingParams(max_tokens=20, temperature=0) + + # Track calls to record_acceptance_rate and the disabled state. + gate_state = {"record_calls": [], "gate_disabled": False} + + original_record_acceptance_rate = SpeculationGate.record_acceptance_rate + + def mock_record_acceptance_rate(self, acceptance_rate, sample_id=None): + """ + Mock that simulates a low true acceptance rate. + This is below the threshold of 0.6, so the gate should trigger after + the window fills. + """ + simulated_low_rate = 0.2 + disabled_now, avg = original_record_acceptance_rate(self, simulated_low_rate, sample_id) + + gate_state["record_calls"].append( + { + "original_rate": acceptance_rate, + "simulated_rate": simulated_low_rate, + "disabled_now": disabled_now, + "avg_acceptance_rate": avg, + "sample_id": sample_id, + } + ) + if disabled_now: + gate_state["gate_disabled"] = True + + return disabled_now, avg + + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + try: + with patch.object(SpeculationGate, "record_acceptance_rate", mock_record_acceptance_rate): + llm_spec.generate(prompts, sampling_params) + + # Verify the mock was called (requests completed) + assert len(gate_state["record_calls"]) > 0, "record_acceptance_rate should have been called" + + # Verify the gate was disabled after enough requests with low acceptance + assert gate_state["gate_disabled"], ( + f"Gate should have been disabled with simulated low acceptance. Calls: {gate_state['record_calls']}" + ) + + # Verify the gate triggered at the right time (after the window is filled). + # The gate should trigger on the `acceptance_rate_window_size`-th call + # (index = window - 1). + disable_indices = [ + i for i, call in enumerate(gate_state["record_calls"]) if call["disabled_now"] + ] + assert len(disable_indices) == 1, ( + f"Gate should have triggered exactly once, but triggered at indices: {disable_indices}" + ) + assert disable_indices[0] >= acceptance_rate_window_size - 1, ( + f"Gate should trigger after window ({acceptance_rate_window_size}) " + f"is filled, but triggered at index {disable_indices[0]}" + ) + + # Verify the average acceptance rate was below threshold when disabled. + disable_call = gate_state["record_calls"][disable_indices[0]] + assert disable_call["avg_acceptance_rate"] is not None + assert disable_call["avg_acceptance_rate"] < acceptance_rate_threshold, ( + f"Avg acceptance rate ({disable_call['avg_acceptance_rate']}) " + f"should be below threshold ({acceptance_rate_threshold})" + ) + + logger.debug(f"Gate correctly triggered after {disable_indices[0] + 1} requests") + logger.debug( + f"Final avg acceptance rate: {disable_call['avg_acceptance_rate']:.3f} " + f"< threshold {acceptance_rate_threshold}" + ) + finally: + llm_spec.shutdown() @pytest.mark.cpu_only diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py b/tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py new file mode 100644 index 000000000000..ebc0e6e91e62 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_user_provided.py @@ -0,0 +1,90 @@ +import unittest + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch.speculative.ngram import NGramDrafter, NGramPoolManager +from tensorrt_llm.llmapi import ( + CudaGraphConfig, + KvCacheConfig, + NGramDecodingConfig, + UserProvidedDecodingConfig, +) + + +# TODO: add disable_overlap_scheduler=False +@pytest.mark.parametrize( + "disable_overlap_scheduler,use_cuda_graph,attn_backend", + [[True, False, "TRTLLM"], [True, True, "TRTLLM"], [True, False, "FLASHINFER"]], +) +def test_llama_user_provided( + disable_overlap_scheduler: bool, use_cuda_graph: bool, attn_backend: str +): + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 20: + pytest.skip("Not enough memory to load target model") + + max_batch_size = 2 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) + cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=llm_models_root() / "llama-3.1-model" / "Meta-Llama-3.1-8B", + backend="pytorch", + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_num_tokens=2048, + ) + + ngram_config = NGramDecodingConfig( + max_draft_len=max_draft_len, + max_matching_ngram_size=2, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=True, + ) + + ngram_pool_manager = NGramPoolManager( + spec_config=ngram_config, + max_num_requests=max_batch_size, + ) + + drafter = NGramDrafter( + spec_config=ngram_config, + ngram_pool_manager=ngram_pool_manager, + ) + + spec_config = UserProvidedDecodingConfig( + max_draft_len=max_draft_len, + drafter=drafter, + ) + + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + sampling_params = SamplingParams(max_tokens=32) + + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(**llm_common_config) + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): + # The spec decode algorithm currently guarantees identical results + assert text_spec == text_ref + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 26379c1b8652..d1b3c613848a 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -14,6 +14,7 @@ # limitations under the License. import json +import os import tempfile import unittest from contextlib import contextmanager @@ -23,8 +24,10 @@ import pytest import torch +from test_common.llm_data import with_mocked_hf_download_for_single_gpu from utils.llm_data import llm_models_root -from utils.util import skip_num_gpus_less_than, skip_pre_blackwell +from utils.util import (skip_blackwell, skip_num_gpus_less_than, + skip_pre_blackwell) from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention.backends.sparse.dsa import ( @@ -32,6 +35,7 @@ from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttentionMetadata from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.pyexecutor._util import ( _derive_draft_max_attention_window, _expand_attention_window_pattern_to_global_layers) @@ -45,6 +49,7 @@ INVALID_PROMPT_LOOKAHEAD_TOKEN from tensorrt_llm._torch.speculative.mtp_dynamic_tree import \ MTPEagleDynamicTreeWorker +from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig) @@ -642,6 +647,280 @@ def test_block_offsets_staging_width_spec_gate(spec_signal): assert draft_kwargs["max_blocks"] is None +@pytest.mark.parametrize( + "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,multi_batch,attention_dp,use_hf_speculative_model", + [ + [True, "TRTLLM", True, False, False, False, False, False, False], + [False, "TRTLLM", True, False, False, False, False, False, False], + [True, "FLASHINFER", True, False, False, False, False, False, False], + [False, "FLASHINFER", True, False, False, False, False, False, False], + [False, "TRTLLM", False, True, True, False, False, False, False], + [True, "TRTLLM", False, True, True, False, False, False, False], + [True, "TRTLLM", True, False, True, True, False, False, False], + [True, "TRTLLM", True, False, True, False, False, False, False], + [True, "TRTLLM", True, False, False, True, False, False, False], + [True, "TRTLLM", False, False, False, False, False, False, False], + [False, "TRTLLM", False, False, False, False, False, False, False], + [True, "TRTLLM", False, False, False, False, True, False, False], + [True, "TRTLLM", False, False, False, False, True, True, False], + [False, "TRTLLM", False, False, False, False, True, False, False], + [True, "TRTLLM", False, False, False, True, False, False, False], + [True, "FLASHINFER", False, False, False, False, False, False, False], + [False, "FLASHINFER", False, False, False, False, False, False, False], + # Tests (mocked) speculative model auto-download from HuggingFace + [False, "TRTLLM", True, False, False, False, False, False, True], + ]) +@pytest.mark.high_cuda_memory +@with_mocked_hf_download_for_single_gpu +def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, + disable_overlap_scheduler: bool, enable_block_reuse: bool, + use_one_model: bool, enable_chunked_prefill: bool, + multi_batch: bool, attention_dp: bool, + use_hf_speculative_model: bool): + if not use_one_model: + pytest.skip("Two model Eagle3 is deprecated") + + # Eagle3 one model works with overlap scheduler and block reuse. + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + if use_hf_speculative_model: + eagle_model = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" + else: + eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + + # bs > 1 gives non-deterministic when doing IFB. There are slight chances + # that ref and spec does not match 100% + max_batch_size = 4 if multi_batch else 1 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, + max_tokens=8192) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[i for i in range(1, max_batch_size + + 1)]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + enable_attention_dp=attention_dp, + # This max_seq_len is larger than the one specified + # in the llama 3 8B eagle's config. We want to make sure + # that the draft model won't go above its max in warmup + # in this test. + max_seq_len=8192, + enable_chunked_prefill=enable_chunked_prefill, + ) + if enable_chunked_prefill: + # Use a small max_num_tokens so that the chunked prefill path gets exercised. + llm_common_config['max_num_tokens'] = 64 + + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=eagle_model, + # Llama 3 does not support one model eagle. + eagle3_one_model=use_one_model, + ) + + # Create the LLM instance + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + # Acceptance rate tests + if enable_chunked_prefill: + # Use a long prompt for chunked prefill tests. + prompts = [ + "The capital of France is a city of romance, art, fashion, and cuisine. Paris is a must-visit destination for anyone who loves history, architecture, and culture. From the iconic Eiffel Tower to the world-famous Louvre Museum, Paris has something to offer for every interest and age.\nThe city is divided into 20 arrondissements, each with its own unique character and charm. The Latin Quarter is a popular area for students and young travelers, while the Champs-ƉlysĆ©es is a hub for shopping and dining. The Montmartre neighborhood is famous for its bohemian vibe and stunning views of the city.\nParis is also known for its beautiful parks and gardens, such as the Luxembourg Gardens and the Tuileries Garden. The city has a rich history, with landmarks like the Notre-Dame Cathedral and the Arc de Triomphe. Visitors can also explore the city's many museums, including the MusĆ©e d'Orsay and the MusĆ©e Rodin.\nIn addition to its cultural and historical attractions, Paris is also a great destination for foodies. The city is famous for its cuisine, including croissants, baguettes, and cheese. Visitors can sample the city's famous dishes at one of the many restaurants, cafes, and " + ] + tok_ids = [llm_spec.tokenizer.encode(prompts[0])] + else: + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + tok_ids = [llm_spec.tokenizer.encode("The future of AI is")] + if multi_batch: + # encode each prompt individually (encode(list) returns nested + # lists in transformers 5.x which prompt_inputs can't handle) + for p in prompts: + tok_ids.append(llm_spec.tokenizer.encode(p)) + + sampling_params = SamplingParams(max_tokens=128, temperature=0) + + for i in range(len(tok_ids)): + num_tokens = 0 + num_drafted = 0 + num_accepted = 0 + + for output in llm_spec.generate_async(tok_ids[i], + sampling_params, + streaming=True): + new_tokens = output.outputs[0].token_ids + num_drafted += max_draft_len + num_accepted += len(new_tokens) - num_tokens - 1 + num_tokens = len(new_tokens) + + accept_rate = num_accepted / num_drafted + assert accept_rate > 0.1 + + # Output tests + sampling_params = SamplingParams(max_tokens=10, temperature=0) + + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(**llm_common_config) + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): + # The spec decode algorithm currently guarantees identical results + assert text_spec == text_ref + + +@pytest.mark.parametrize("eagle3_one_model", [True, False]) +def test_eagle3_spec_decoding_stats(eagle3_one_model): + """Test that specDecodingStats are correctly populated in metrics endpoint""" + models_path = llm_models_root() + eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + # Skip if models don't exist + if not os.path.exists(target_model_dir) or not os.path.exists( + eagle_model_dir): + pytest.skip(f"Required models not found") + + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.6) + spec_config = Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=eagle_model_dir, + eagle3_one_model=eagle3_one_model, + ) + + with LLM( + model=target_model_dir, + speculative_config=spec_config, + kv_cache_config=kv_cache_config, + disable_overlap_scheduler=not eagle3_one_model, + enable_iter_perf_stats=True, + max_batch_size=4, + ) as llm: + # Generate some output to collect stats + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + sampling_params = SamplingParams(max_tokens=20, temperature=0) + llm.generate(prompts, sampling_params) + + # Get iteration stats + stats = llm.get_stats(timeout=2) + assert len(stats) > 0, "Should have iteration stats" + + # Find iterations with speculation (generation phase) + iterations_with_spec = [] + for stat in stats: + if 'specDecodingStats' in stat: + spec_stats = stat['specDecodingStats'] + if spec_stats.get('numDraftTokens', 0) > 0: + iterations_with_spec.append(spec_stats) + + # Should have at least some iterations with spec decoding + assert len(iterations_with_spec) > 0, \ + f"Should have iterations with specDecodingStats (found {len(iterations_with_spec)})" + + # Validate specDecodingStats structure and values + for spec_stats in iterations_with_spec: + # Check all fields are present + assert 'numDraftTokens' in spec_stats + assert 'numAcceptedTokens' in spec_stats + assert 'numRequestsWithDraftTokens' in spec_stats + assert 'acceptanceLength' in spec_stats + assert 'iterLatencyMS' in spec_stats + assert 'draftOverhead' in spec_stats + + # Validate value constraints + assert spec_stats['numDraftTokens'] > 0 + assert 0 <= spec_stats['numAcceptedTokens'] <= spec_stats[ + 'numDraftTokens'] + assert spec_stats['numRequestsWithDraftTokens'] > 0 + assert spec_stats['acceptanceLength'] >= 1.0 + assert spec_stats['iterLatencyMS'] >= 0.0 + assert 0.0 <= spec_stats['draftOverhead'] <= 1.0 + + # Calculate overall acceptance rate + total_draft = sum(s['numDraftTokens'] for s in iterations_with_spec) + total_accepted = sum(s['numAcceptedTokens'] + for s in iterations_with_spec) + acceptance_rate = (total_accepted / total_draft * + 100) if total_draft > 0 else 0 + + # Should have reasonable acceptance rate for Eagle3 + assert acceptance_rate > 5.0, f"Acceptance rate too low: {acceptance_rate:.1f}%" + + +@pytest.mark.parametrize("use_cuda_graph", [True, False]) +@pytest.mark.high_cuda_memory +def test_llama_eagle3_long_prompt(use_cuda_graph): + # Eagle3 one model works with overlap scheduler and block reuse. + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + spec_config = Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=eagle_model_dir, + eagle3_one_model=False, + ) + + if use_cuda_graph: + cuda_graph_config = CudaGraphConfig(batch_sizes=[1]) + else: + cuda_graph_config = None + + llm_spec = LLM(model=target_model_dir, + speculative_config=spec_config, + max_batch_size=1, + cuda_graph_config=cuda_graph_config, + disable_overlap_scheduler=True) + + prompt = [", ".join(str(i) for i in range(1000))] + + sampling_params = SamplingParams(max_tokens=10, temperature=0) + results_spec = llm_spec.generate(prompt, sampling_params) + + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(model=target_model_dir, + max_batch_size=1, + cuda_graph_config=None, + disable_overlap_scheduler=False) + + results_ref = llm_ref.generate(prompt, sampling_params) + + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + # The LLM with speculation on should dynamically turn it off in this + # test since it goes beyond the max seqlen. Thus, the text should be + # _exactly_ the same, no need to use similarity scoring. + assert generated_text_spec[0] == generated_text_ref[0] + + def test_deepseek_mla_eagle3(): use_cuda_graph = True attn_backend = "TRTLLM" @@ -750,6 +1029,172 @@ def test_deepseek_mla_eagle3(): pass +@pytest.mark.parametrize("use_one_model", [True, False]) +def test_multi_eagle3(use_one_model: bool): + use_cuda_graph = True + attn_backend = "TRTLLM" + disable_overlap_scheduler = False + enable_block_reuse = False + enable_chunked_prefill = False + + # Eagle3 one model works with overlap scheduler and block reuse. + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 150: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + eagle_config = { + 'architectures': ['LlamaForCausalLMEagle3'], + 'attention_bias': False, + 'attention_dropout': 0.0, + 'bos_token_id': 128000, + 'eos_token_id': [128001, 128008, 128009], + 'eagle_config': { + 'use_aux_hidden_state': False, + 'use_input_layernorm_in_first_layer': True, + 'use_last_layernorm': True, + 'use_mtp_layernorm': False + }, + 'head_dim': 128, + 'hidden_act': 'silu', + 'hidden_size': 4096, + 'initializer_range': 0.02, + 'intermediate_size': 16384, + 'max_position_embeddings': 131072, + 'mlp_bias': False, + 'model_type': 'llama', + 'num_attention_heads': 32, + 'num_eagle_features': 1, + 'num_hidden_layers': 2, + 'num_key_value_heads': 8, + 'pretraining_tp': 1, + 'rms_norm_eps': 1e-05, + 'rope_scaling': { + 'factor': 8.0, + 'high_freq_factor': 4.0, + 'low_freq_factor': 1.0, + 'original_max_position_embeddings': 8192, + 'rope_type': 'llama3' + }, + 'rope_theta': 500000.0, + 'tie_word_embeddings': False, + 'torch_dtype': 'bfloat16', + 'transformers_version': '4.52.4', + 'use_cache': True, + 'vocab_size': 128256, + 'draft_vocab_size': 128256, + } + with tempfile.TemporaryDirectory() as temp_dir: + eagle_model_dir = Path(temp_dir) + config_path = eagle_model_dir / "config.json" + with config_path.open("w") as f: + json.dump(eagle_config, f, indent=2) + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + + # bs > 1 gives non-deterministic when doing IFB. There are slight chances + # that ref and spec does not match 100% + max_batch_size = 16 + max_draft_len = 3 + kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, + free_gpu_memory_fraction=0.5) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + attn_backend=attn_backend, + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + enable_chunked_prefill=enable_chunked_prefill, + load_format="dummy", + ) + + spec_config = Eagle3DecodingConfig(max_draft_len=max_draft_len, + speculative_model=eagle_model_dir, + eagle3_one_model=use_one_model, + load_format="dummy") + + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + tok_ids = llm_spec.tokenizer.encode("The future of AI is") + + sampling_params = SamplingParams(max_tokens=32, temperature=0) + for output in llm_spec.generate_async(tok_ids, + sampling_params, + streaming=True): + pass + + +@pytest.mark.parametrize("use_dynamic_tree", [False, True], + ids=["no_dynamic_tree", "dynamic_tree"]) +@pytest.mark.parametrize("use_cuda_graph", [False, True]) +@pytest.mark.high_cuda_memory +@skip_blackwell +# Opt out of MPI session reuse: the XQA JIT cubin registry is process-global +# (DecoderXQARunner::getResourceGlobal) and its lookup key does not include +# q_seq_len / is_spec_dec_tree, so running the dynamic-tree and non-dynamic-tree +# variants in one worker process launches a cubin compiled for the other +# config's q_seq_len -> CUDA_ERROR_INVALID_VALUE on Hopper. +@pytest.mark.private_mpi_session +@with_mocked_hf_download_for_single_gpu +def test_llama_eagle3_rejection_sampling_modes(use_dynamic_tree: bool, + use_cuda_graph: bool): + """Test one-model rejection sampling with and without dynamic tree.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + + max_batch_size = 1 + max_draft_len = 6 + dynamic_tree_max_top_k = 10 + max_total_draft_tokens = 60 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, max_tokens=8192) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=True, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=8192, + ) + + spec_config_kwargs = dict( + max_draft_len=max_draft_len, + speculative_model=eagle_model, + eagle3_one_model=True, + use_rejection_sampling=True, + ) + if use_dynamic_tree: + spec_config_kwargs.update( + use_dynamic_tree=True, + dynamic_tree_max_topK=dynamic_tree_max_top_k, + max_total_draft_tokens=max_total_draft_tokens, + ) + + llm_spec = LLM(**llm_common_config, + speculative_config=Eagle3DecodingConfig( + **spec_config_kwargs)) + + prompts = ["The president of the United States is"] + sampling_params = SamplingParams(max_tokens=20, temperature=1.0, top_p=1.0) + + results = llm_spec.generate(prompts, sampling_params) + llm_spec.shutdown() + + assert len(results) == len(prompts) + assert len(results[0].outputs[0].token_ids) > 0 + + @pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) @pytest.mark.parametrize("use_cuda_graph", [False, True]) @pytest.mark.high_cuda_memory @@ -828,5 +1273,167 @@ def test_nemotron_super_mtp_dynamic_tree_dl6_k10_dt31( assert text_spec == text_ref +@pytest.mark.parametrize("use_cuda_graph", [True, False]) +def test_eagle3_lora(use_cuda_graph: bool): + """Test LoRA with 3 requests and max_batch_size=4. + + This test verifies that when using LoRA modules, + the system properly applies the LoRA configurations. + """ + attn_backend = "TRTLLM" + enable_block_reuse = False + use_one_model = True + enable_chunked_prefill = False + + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + + eagle_model_dir = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + hf_lora_dir = f"{models_path}/llama-models/luotuo-lora-7b-0.1" + + # Test with 3 requests and max_batch_size=4 to trigger padding + max_batch_size = 4 + max_draft_len = 4 + kv_cache_config = KvCacheConfig(enable_block_reuse=enable_block_reuse, + max_tokens=8192) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[1, 2, 4], enable_padding=True) if use_cuda_graph else None + lora_config = LoraConfig(max_lora_rank=64, max_loras=2, max_cpu_loras=2) + + llm_common_config = dict( + model=target_model_dir, + attn_backend=attn_backend, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=1024, + enable_chunked_prefill=enable_chunked_prefill, + lora_config=lora_config, + ) + + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=eagle_model_dir, + eagle3_one_model=use_one_model, + ) + + # Create the LLM instance + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + prompts = [ + "The capital of France is", + "The president of the United States is", + "The future of AI is", + ] + lora_requests = [LoRARequest("luotuo", 1, hf_lora_dir)] * len(prompts) + + sampling_params = SamplingParams(max_tokens=20, temperature=0) + llm_spec.generate(prompts, sampling_params, lora_request=lora_requests) + llm_spec.shutdown() + + +@pytest.mark.parametrize("disable_overlap_scheduler", [False]) +@pytest.mark.parametrize("use_cuda_graph", [True]) +@pytest.mark.high_cuda_memory +@with_mocked_hf_download_for_single_gpu +def test_llama_eagle3_dynamic_tree(use_cuda_graph: bool, + disable_overlap_scheduler: bool): + """Test EAGLE3 dynamic tree speculative decoding with one-model architecture.""" + total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 + if total_mem_gb < 35: + pytest.skip("Not enough memory to load target + draft model") + + models_path = llm_models_root() + target_model_dir = f"{models_path}/llama-3.1-model/Llama-3.1-8B-Instruct" + eagle_model = f"{models_path}/EAGLE3-LLaMA3.1-Instruct-8B" + + max_batch_size = 4 + max_draft_len = 6 + dynamic_tree_max_topK = 10 + max_total_draft_tokens = 30 + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + max_tokens=2048, + free_gpu_memory_fraction=0.5) + cuda_graph_config = CudaGraphConfig( + batch_sizes=[i for i in range(1, max_batch_size + + 1)]) if use_cuda_graph else None + + llm_common_config = dict( + model=target_model_dir, + attn_backend="TRTLLM", + disable_overlap_scheduler=disable_overlap_scheduler, + cuda_graph_config=cuda_graph_config, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + max_seq_len=2048, + ) + + spec_config = Eagle3DecodingConfig( + max_draft_len=max_draft_len, + speculative_model=eagle_model, + eagle3_one_model=True, + use_dynamic_tree=True, + dynamic_tree_max_topK=dynamic_tree_max_topK, + max_total_draft_tokens=max_total_draft_tokens, + ) + + # Create the LLM instance + llm_spec = LLM(**llm_common_config, speculative_config=spec_config) + + # Acceptance rate tests + prompts = [ + "The capital of France is", + "The president of the United States is", + ] + tok_ids = [llm_spec.tokenizer.encode("The future of AI is")] + + sampling_params = SamplingParams(max_tokens=128, temperature=0) + + for i in range(len(tok_ids)): + num_tokens = 0 + num_drafted = 0 + num_accepted = 0 + + for output in llm_spec.generate_async(tok_ids[i], + sampling_params, + streaming=True): + new_tokens = output.outputs[0].token_ids + num_drafted += max_draft_len + num_accepted += len(new_tokens) - num_tokens - 1 + num_tokens = len(new_tokens) + + accept_rate = num_accepted / num_drafted + assert accept_rate > 0.10 + + # Output tests: verify spec decode matches reference + sampling_params = SamplingParams(max_tokens=10, temperature=0) + + results_spec = llm_spec.generate(prompts, sampling_params) + generated_text_spec = [result.outputs[0].text for result in results_spec] + llm_spec.shutdown() + + llm_ref = LLM(**llm_common_config) + results_ref = llm_ref.generate(prompts, sampling_params) + generated_text_ref = [result.outputs[0].text for result in results_ref] + llm_ref.shutdown() + + def assert_meaningful_text(text: str) -> None: + stripped = text.strip() + assert stripped + assert "\ufffd" not in stripped + assert any(ch.isalpha() for ch in stripped) + words = stripped.lower().split() + assert not any( + len(set(words[i:i + 6])) == 1 for i in range(len(words) - 5)) + + for text_spec, text_ref in zip(generated_text_spec, generated_text_ref): + assert_meaningful_text(text_spec) + assert_meaningful_text(text_ref) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py b/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py index 899c6810be85..9b726263e174 100644 --- a/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py +++ b/tests/unittest/auto_deploy/_utils_test/_model_test_utils.py @@ -428,6 +428,15 @@ def apply_rotary_pos_emb_ds(q, k, cos, sin, position_ids, unsqueeze_dim=1): _SMALL_MODEL_CONFIGS = { + "meta-llama/Meta-Llama-3.1-8B-Instruct": { + "model_kwargs": { + "num_hidden_layers": 1, + "hidden_size": 64, + "intermediate_size": 64, + "num_attention_heads": 2, + "num_key_value_heads": 1, + }, + }, "Qwen/Qwen3-30B-A3B": { "model_kwargs": { "num_hidden_layers": 2, @@ -562,7 +571,7 @@ def get_small_model_config(model_hub_id: str, **llm_args_kwargs) -> Dict[str, An Get the small model configuration for a given HuggingFace model hub ID. Args: - model_hub_id: The HuggingFace model hub ID (e.g., "Qwen/Qwen3-30B-A3B") + model_hub_id: The HuggingFace model hub ID (e.g., "meta-llama/Meta-Llama-3.1-8B-Instruct") Returns: Dictionary containing the model configuration diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py new file mode 100644 index 000000000000..2517196708e2 --- /dev/null +++ b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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 signal +import subprocess +import tempfile +from contextlib import contextmanager +from pathlib import Path + +import pytest +import torch +import yaml +from _model_test_utils import get_small_model_config +from click.testing import CliRunner +from utils.cpp_paths import llm_root # noqa: F401 + +from tensorrt_llm.commands.bench import main + +# needed since LLM API uses MPI executor pool internally for TP>1, which leaks a thread on shutdown +pytestmark = pytest.mark.threadleak(enabled=False) + + +class TimeoutError(Exception): + """Exception raised when a test times out.""" + + pass + + +@contextmanager +def timeout(seconds): + """Context manager that raises TimeoutError if code block exceeds time limit. + + Args: + seconds: Maximum time in seconds to allow the code block to run + + Raises: + TimeoutError: If the code block execution exceeds the time limit + """ + + def timeout_handler(signum, frame): + raise TimeoutError(f"Test execution exceeded {seconds} seconds timeout") + + # Set the signal handler and alarm + old_handler = signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(seconds) + try: + yield + finally: + # Restore the old signal handler and cancel the alarm + signal.alarm(0) + signal.signal(signal.SIGALRM, old_handler) + + +@pytest.fixture(scope="module", autouse=True) +def prewarm_flashinfer_jit(): + """Pre-warm FlashInfer JIT kernels before multi-GPU tests. + + This prevents a race condition where multiple MPI ranks try to JIT-compile + FlashInfer kernels simultaneously to the same cache directory, causing + Ninja build failures like: "ninja: error: opening build log: No such file or directory" + + By triggering the compilation in the main process first, the kernels are + cached and available for all worker ranks. + """ + try: + import flashinfer + import flashinfer.page + import flashinfer.sampling + + if torch.cuda.is_available(): + # Prevent concurrent JIT warmup across multiple pytest processes (e.g., xdist). + try: + import fcntl # Linux-only + except ImportError: + fcntl = None + + lock_f = None + if fcntl is not None: + import pathlib + import tempfile + + lock_path = pathlib.Path(tempfile.gettempdir()) / "flashinfer_jit_prewarm.lock" + lock_f = open(lock_path, "w") + fcntl.flock(lock_f.fileno(), fcntl.LOCK_EX) + # Create dummy tensors to trigger kernel JIT compilation + with torch.no_grad(): + device = torch.device("cuda:0") + + # Trigger page kernel compilation + try: + # Force module loading (this triggers JIT compilation) + _ = flashinfer.page.gen_page_module() + except Exception as exc: # noqa: BLE001 + import warnings + + warnings.warn(f"FlashInfer page-kernel prewarm failed: {exc!r}", RuntimeWarning) + + # Trigger sampling kernel compilation + try: + dummy_probs = torch.softmax(torch.randn(1, 100, device=device), dim=-1) + _ = flashinfer.sampling.sampling_from_probs(dummy_probs, deterministic=True) + except Exception as exc: # noqa: BLE001 + import warnings + + warnings.warn( + f"FlashInfer sampling-kernel prewarm failed: {exc!r}", RuntimeWarning + ) + + torch.cuda.empty_cache() + if lock_f is not None: + lock_f.close() + + except ImportError: + pass # FlashInfer not available + + yield + + +@pytest.fixture(scope="module") +def shared_dataset(llm_root): # noqa: F811 + """Prepare dataset once for all tests in this module.""" + model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct" + config = get_small_model_config(model_name) + with tempfile.TemporaryDirectory() as temp_dir: + dataset_path = _prepare_dataset( + llm_root, temp_dir, config["args"]["model"], num_requests=10 + ) + # Read dataset content to return it (temp_dir will be deleted) + with open(dataset_path, "r") as f: + dataset_content = f.read() + yield dataset_content + + +def _prepare_dataset(root_dir: str, temp_dir: str, model_path_or_name: str, num_requests: int = 10): + """Prepare a synthetic dataset for benchmarking.""" + _DATASET_NAME = "synthetic_128_128.txt" + dataset_path = Path(temp_dir, _DATASET_NAME) + + # Generate a small dataset to run a test - matching workload configuration. + # Write straight to a file via --output (the documented usage) instead of + # scraping --stdout: trtllm-bench prints an import-time banner on stdout that + # would otherwise corrupt the captured JSONL. + command = [ + "trtllm-bench", + "--model", + model_path_or_name, + "prepare-dataset", + "--output", + str(dataset_path), + "token-norm-dist", + "--input-mean", + "128", + "--output-mean", + "128", + "--input-stdev", + "0", + "--output-stdev", + "0", + "--num-requests", + str(num_requests), + ] + print(f"Running command: {' '.join(command)}") + result = subprocess.run(command, cwd=str(temp_dir), capture_output=True, text=True, timeout=300) + if result.returncode != 0: + raise RuntimeError(f"Failed to prepare dataset: {result.stderr}") + return dataset_path + + +@pytest.mark.parametrize( + "allreduce_strategy", + [ + "AUTO", + "ONESHOT", + "TWOSHOT", + "MIN_LATENCY", + "NCCL", + "SYMM_MEM", + ], +) +def test_allreduce_strategies(llm_root, shared_dataset, allreduce_strategy): # noqa: F811 + """Test different allreduce strategies with multi-GPU configuration making sure that there are no crashes or hangs. + + Configuration: + The allreduce_strategy is set in the transforms config: + ```yaml + transforms: + detect_sharding: + allreduce_strategy: "ONESHOT" # or AUTO, NCCL, TWOSHOT, etc. + ``` + + Test configuration: + - Model: Llama-3.1-8B with TP=2 + - Dataset: 10 synthetic requests (128 input, 128 output tokens) + - Timeout: 300 seconds to catch hangs + - Skipped if fewer than 2 GPUs available + + Args: + llm_root: Root directory fixture + shared_dataset: Shared dataset fixture (prepared once for all test runs) + allreduce_strategy: Strategy to test (AUTO, ONESHOT, TWOSHOT, MIN_LATENCY, NCCL) + """ + # Fixed timeout for all strategies (5 minutes should be enough) + TEST_TIMEOUT_SECONDS = 300 + + model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct" + # Override hidden_size to a multiple of one warp's worth of 128-bit accesses (32 threads * + # 8 fp16 elements/access = 256) so the fused all-reduce/RMSNorm kernel never launches a + # partial-warp block. This isolates whether CI failures are specific to the partial-warp + # code path or are unrelated infra flakiness that also affects the full-warp path. + config = get_small_model_config( + model_name, + model_kwargs={ + "num_hidden_layers": 1, + "hidden_size": 256, + "intermediate_size": 256, + "num_attention_heads": 2, + "num_key_value_heads": 1, + }, + ) + tp_size = 2 + max_batch_size = 256 + max_num_tokens = 8192 + + if not torch.cuda.is_available() or torch.cuda.device_count() < tp_size: + pytest.skip(f"Allreduce strategy test requires at least {tp_size} GPUs, skipping") + + with tempfile.TemporaryDirectory() as temp_dir: + # Write shared dataset to temp location + dataset_path = Path(temp_dir, "synthetic_128_128.txt") + with open(dataset_path, "w") as f: + f.write(shared_dataset) + + # Create configuration with specified allreduce strategy in transforms + extra_llm_api_options_path = f"{temp_dir}/extra_llm_api_options.yaml" + with open(extra_llm_api_options_path, "w") as f: + yaml.dump( + { + **config["args"], + "max_batch_size": max_batch_size, + "max_num_tokens": max_num_tokens, + "max_seq_len": 256, + "transforms": { + "detect_sharding": { + "stage": "sharding", + "allreduce_strategy": allreduce_strategy, + }, + "compile_model": { + "stage": "compile", + "backend": "torch-cudagraph", + "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128, 256], + }, + }, + }, + f, + ) + + # Run benchmark with specified allreduce strategy with timeout protection + runner = CliRunner() + args = [ + "--model", + model_name, + ] + + # Only pass --model_path if it's a local filesystem path + # Note: --model_path must come BEFORE the subcommand (throughput) + if str(config["args"]["model"]).startswith("/"): + args.extend(["--model_path", str(config["args"]["model"])]) + + # Add the subcommand and its options + args.extend( + [ + "throughput", + "--backend", + "_autodeploy", + "--dataset", + str(dataset_path), + "--extra_llm_api_options", + extra_llm_api_options_path, + "--tp", + str(tp_size), + "--max_batch_size", + str(max_batch_size), + "--max_num_tokens", + str(max_num_tokens), + ] + ) + + try: + with timeout(TEST_TIMEOUT_SECONDS): + result = runner.invoke(main, args, catch_exceptions=False) + assert result.exit_code == 0, f"Benchmark failed with output: {result.output}" + except TimeoutError as e: + pytest.fail( + f"Test timed out after {TEST_TIMEOUT_SECONDS}s for strategy {allreduce_strategy}. " + f"This might indicate a hang (e.g., TWOSHOT without C++ fix). Error: {e}" + ) diff --git a/tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py new file mode 100644 index 000000000000..08c449198474 --- /dev/null +++ b/tests/unittest/auto_deploy/multigpu/smoke/test_ad_build_small_multi.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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. +"""Testing build_and_run_ad end2end.""" + +import pytest + +__extra_import_path__ = ["~/examples/auto_deploy"] +from _model_test_utils import get_small_model_config +from build_and_run_ad import ExperimentConfig, main + + +@pytest.mark.parametrize("world_size", [1, 2]) +@pytest.mark.parametrize( + "model_hub_id, llm_extra_args", + [ + ( + "meta-llama/Meta-Llama-3.1-8B-Instruct", + { + "transforms": { + "insert_cached_attention": {"backend": "flashinfer"}, + # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/9878 + # "compile_model": {"backend": "torch-opt"}, + "compile_model": {"backend": "torch-cudagraph"}, + }, + }, + ), + ], +) +def test_build_ad(world_size: int, model_hub_id: str, llm_extra_args: dict): + experiment_config = get_small_model_config(model_hub_id, **llm_extra_args) + + experiment_config["args"]["world_size"] = world_size + experiment_config["args"]["runtime"] = "trtllm" # Default runtime set to trtllm + + experiment_config = ExperimentConfig(**experiment_config) + print(f"Experiment Config: {experiment_config}") + main(experiment_config) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py index d3f925185d97..a146136bf820 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_eagle.py @@ -16,6 +16,7 @@ """Unit tests for Eagle3 model with AutoDeploy.""" from pathlib import Path +from typing import Any, ClassVar, Dict from unittest.mock import patch import pytest @@ -24,8 +25,15 @@ import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm -from tensorrt_llm._torch.auto_deploy.models.custom.modeling_eagle import EagleRMSNorm, EagleWrapper +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_eagle import ( + Eagle3DraftOutput, + EagleConfig, + EagleDrafterForCausalLM, + EagleRMSNorm, + EagleWrapper, +) from tensorrt_llm._torch.auto_deploy.models.eagle import EagleDrafterFactory +from tensorrt_llm._torch.auto_deploy.models.factory import ModelFactoryRegistry from tensorrt_llm._torch.auto_deploy.utils.node_utils import ( get_weight_shape, infer_draft_embedding_size, @@ -33,6 +41,7 @@ ) __extra_import_path__ = ["~/examples/auto_deploy"] +from build_and_run_ad import ExperimentConfig, main from test_common.llm_data import hf_id_to_local_model_dir EAGLE_MODEL_HUB_ID = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" @@ -59,6 +68,98 @@ "topk_group": 1, } +############################################################################### +# Mock classes for standalone Eagle testing +# +# These classes enable unit testing the Eagle checkpoint without a target model. +# In production speculative decoding, real hidden states come from the target model. +# For testing, MockEagle3ModelForCausalLM generates random hidden states. +############################################################################### + + +class MockEagleConfig(EagleConfig): + """Config for standalone Eagle testing with embedding/lm_head loaded from checkpoint. + + In production, Eagle shares embedding/lm_head with the target model. + For standalone testing, we need to load these from the checkpoint. + """ + + _drafter_defaults: ClassVar[Dict[str, Dict[str, Any]]] = { + "llama": { + "load_embedding_from_target": False, + "load_lm_head_from_target": False, + "num_capture_layers": 1, + }, + } + + +class MockEagle3ModelForCausalLM(EagleDrafterForCausalLM): + """Test wrapper that provides random hidden states for standalone Eagle testing. + + In production speculative decoding, real hidden states come from the target model. + This mock class generates random hidden states for testing the Eagle model in isolation. + """ + + def __init__(self, config): + super().__init__(config) + self._hidden_size = config.hidden_size + self._dtype = config.dtype + + def forward(self, input_ids, position_ids, input_embeds=None, **kwargs): + assert self.model.embed_tokens is not None, ( + "embed_tokens must be set before running standalone Eagle model." + ) + assert self.lm_head is not None, ( + "lm_head must be set before running standalone Eagle model." + ) + + if input_embeds is None: + inputs_embeds = self.model.embed_tokens(input_ids) + + # Inject mock hidden states if not provided + if "hidden_states" not in kwargs: + batch_size, seq_len = input_ids.shape + kwargs["hidden_states"] = torch.randn( + (batch_size, seq_len, self._hidden_size), + dtype=self._dtype, + device=input_ids.device, + ) + draft_output = super().forward(inputs_embeds, position_ids, **kwargs) + logits = self.lm_head(draft_output.norm_hidden_state) + return Eagle3DraftOutput(logits=logits, last_hidden_state=draft_output.last_hidden_state) + + +class MockEagleDrafterFactory(EagleDrafterFactory): + """Test factory that uses MockEagle3ModelForCausalLM for standalone Eagle testing. + + This factory directly builds MockEagle3ModelForCausalLM with MockEagleConfig, + which loads embedding/lm_head from checkpoint for standalone testing. + """ + + def _build_model(self, device): + from contextlib import nullcontext + + from accelerate import init_empty_weights + + model_config, unused_kwargs = self._get_model_config() + # transformers>=5.5 applies @dataclass(kw_only=True) to PretrainedConfig + # subclasses, overriding EagleConfig.__init__. Use the factory classmethod. + model_config = MockEagleConfig.from_base_config(model_config, model_config.model_type) + + with (init_empty_weights if device == "meta" else nullcontext)(): + model = MockEagle3ModelForCausalLM._from_config(model_config, **unused_kwargs) + + if device == "meta": + if hasattr(model, "post_init"): + model.post_init() + else: + model.to(device) + + self._checkpoint_conversion_mapping = getattr(model, "_checkpoint_conversion_mapping", None) + model.eval() + + return model + def _build_small_draft_factory( model_hub_id: str, model_kwargs: dict | None = None @@ -84,6 +185,48 @@ def test_eagle_rmsnorm_keeps_fp32_weights(): assert norm.weight.dtype == torch.float32 +@pytest.fixture +def register_mock_eagle_factory(): + """Register MockEagleDrafterFactory for the test and clean up afterwards. + + This fixture temporarily registers the mock factory with ModelFactoryRegistry, + allowing tests to use model_factory="MockEagleDrafter", and removes the + registration after the test completes. + """ + ModelFactoryRegistry._registry["MockEagleDrafter"] = MockEagleDrafterFactory + yield + ModelFactoryRegistry._registry.pop("MockEagleDrafter", None) + + +def test_build_ad_eagle(register_mock_eagle_factory): + """Test building Eagle model with AutoDeploy using MockEagleDrafterFactory. + + This test uses the MockEagleDrafterFactory which builds MockEagle3ModelForCausalLM, + a mock model that generates random hidden states for standalone Eagle testing. + """ + llm_extra_args = { + "model_factory": "MockEagleDrafter", + "transforms": { + "insert_cached_attention": {"backend": "trtllm"}, + "compile_model": { + "backend": "torch-simple", + "piecewise_enabled": False, + }, + }, + } + experiment_config = get_small_model_config(EAGLE_MODEL_HUB_ID, **llm_extra_args) + experiment_config["args"]["runtime"] = "demollm" + experiment_config["args"]["world_size"] = 0 + experiment_config["args"]["tokenizer"] = hf_id_to_local_model_dir( + "meta-llama/Meta-Llama-3.1-8B-Instruct" + ) + + print(f"Experiment Config: {experiment_config}") + experiment_config = ExperimentConfig(**experiment_config) + + main(experiment_config) + + def test_eagle_model_torch_export(): """Test that Eagle3Model can be exported with torch.export. diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py index bcd24492a53f..b05dc5949634 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_build_small_single.py @@ -84,6 +84,32 @@ def _check_ad_config(experiment_config: ExperimentConfig, llm_args: LlmArgs): @pytest.mark.parametrize( "model_hub_id, llm_extra_args", [ + ( + "meta-llama/Meta-Llama-3.1-8B-Instruct", + { + "kv_cache_config": { + "free_gpu_memory_fraction": 0.0001, + }, + "transforms": { + "insert_cached_attention": {"backend": "flashinfer"}, + # TODO: https://github.com/NVIDIA/TensorRT-LLM/issues/9878 + # "compile_model": {"backend": "torch-opt"}, + "compile_model": { + "backend": "torch-cudagraph", + "cuda_graph_batch_sizes": [1, 2], + }, + }, + }, + ), + ( + "meta-llama/Meta-Llama-3.1-8B-Instruct", + { + "transforms": { + "transformers_replace_cached_attn": {"backend": "flashinfer"}, + }, + "mode": "transformers", + }, + ), ( "Qwen/Qwen3-30B-A3B", { diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py new file mode 100644 index 000000000000..2ee4ce344b0e --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_guided_decoding_regex.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +from _model_test_utils import get_small_model_config + +from tensorrt_llm.llmapi import GuidedDecodingParams + +__extra_import_path__ = ["~/examples/auto_deploy"] +from build_and_run_ad import ExperimentConfig, main + + +def test_ad_guided_decoding_regex_e2e(): + """Test guided decoding with regex pattern validation using the build_and_run_ad main().""" + test_case = { + "prompt": "What is the capital of France?", + "regex": r"I don't know, I am a randomly initialized model|Paris", + "valid_responses": ["I don't know, I am a randomly initialized model", "Paris"], + } + + guided_decoding_backend = "xgrammar" + + experiment_config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") + + # DemoLLM runtime does not support guided decoding. Need to set runtime to trtllm. + experiment_config["args"]["runtime"] = "trtllm" + experiment_config["args"]["world_size"] = 1 + # NOTE: trtllm attention backend fails on B200 (likely illegal memory access); use flashinfer. + experiment_config["args"]["attn_backend"] = "flashinfer" + experiment_config["args"]["guided_decoding_backend"] = guided_decoding_backend + + experiment_config["prompt"]["batch_size"] = 1 + experiment_config["prompt"]["queries"] = test_case["prompt"] + + cfg = ExperimentConfig(**experiment_config) + + # Need to introduce the guided decoding params after ExperimentConfig construction + # because otherwise they get unpacked as a dict. + cfg.prompt.sp_kwargs = { + "max_tokens": 10, + "top_k": None, + "temperature": 0.1, + "guided_decoding": GuidedDecodingParams(regex=test_case["regex"]), + } + + print(f"Experiment config: {experiment_config}") + print("Generating outputs...") + results = main(cfg) + print("Results:", results) + + # Parse and validate: output should be a prefix of one of the valid responses + prompts_and_outputs = results["prompts_and_outputs"] + assert len(prompts_and_outputs) == 1 + generated_text = prompts_and_outputs[0][1].strip() + + valid_responses = test_case["valid_responses"] + is_valid_prefix = any(response.startswith(generated_text) for response in valid_responses) + assert is_valid_prefix, ( + f"Output is not a valid prefix of any expected response.\n" + f"Generated: '{generated_text}'\n" + f"Valid responses: {valid_responses}" + ) diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py index 2c5d3ffe5a00..849e7e8eefa9 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_speculative_decoding.py @@ -14,15 +14,43 @@ # limitations under the License. +import torch from _model_test_utils import get_small_model_config -from tensorrt_llm.llmapi import MTPDecodingConfig +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models.eagle import EagleOneModelFactory +from tensorrt_llm._torch.auto_deploy.transform.interface import TransformConfig +from tensorrt_llm._torch.auto_deploy.transform.library.hidden_states import ( + DetectHiddenStatesForCapture, +) +from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens +from tensorrt_llm.llmapi import Eagle3DecodingConfig, MTPDecodingConfig __extra_import_path__ = ["~/examples/auto_deploy"] from build_and_run_ad import ExperimentConfig, main from test_common.llm_data import hf_id_to_local_model_dir +def get_extra_seq_len_for_kv_cache(llm_args) -> int: + """Mirror the current extra-KV sizing logic used by the runtime.""" + extra = 0 + spec_config = llm_args.speculative_config + if not llm_args.disable_overlap_scheduler: + extra += 1 + if spec_config is not None: + extra += spec_config.tokens_per_gen_step - 1 + + if spec_config is not None: + extra += spec_config.tokens_per_gen_step - 1 + extra += get_num_extra_kv_tokens(spec_config) + + return extra + + +def piecewise_disabled_transforms(): + return {"compile_model": {"piecewise_enabled": False}} + + def test_super_mtp_smoke(): """Test one-model MTP/Eagle runtime with a tiny Nemotron SuperV3 target.""" test_prompt = "What is the capital of France?" @@ -138,3 +166,129 @@ def test_super_mtp_ssm_replay_smoke(): prompts_and_outputs = results["prompts_and_outputs"] assert len(prompts_and_outputs) == 1 + + +def test_kv_cache_extra_seq_len_for_spec_dec(): + """Test that get_extra_seq_len_for_kv_cache computes correct extra capacity.""" + from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs + + # Case 1: No spec config, no overlap + args_no_spec = LlmArgs( + model="meta-llama/Meta-Llama-3.1-8B-Instruct", + disable_overlap_scheduler=True, + ) + assert get_extra_seq_len_for_kv_cache(args_no_spec) == 0 + + # Case 2: No spec config, with overlap + args_overlap = LlmArgs( + model="meta-llama/Meta-Llama-3.1-8B-Instruct", + disable_overlap_scheduler=False, + ) + assert get_extra_seq_len_for_kv_cache(args_overlap) == 1 # overlap adds +1 + + # Case 3: Eagle3 one-model, overlap disabled + spec_config = Eagle3DecodingConfig( + max_draft_len=3, + speculative_model="some/model", + eagle3_one_model=True, + ) + args_eagle = LlmArgs( + model="meta-llama/Meta-Llama-3.1-8B-Instruct", + speculative_config=spec_config, + disable_overlap_scheduler=True, + transforms=piecewise_disabled_transforms(), + ) + extra = get_extra_seq_len_for_kv_cache(args_eagle) + # Should include max_total_draft_tokens + get_num_extra_kv_tokens (max_draft_len - 1) + assert extra > 0 + assert extra == spec_config.max_total_draft_tokens + (spec_config.max_draft_len - 1) + + # Case 4: Eagle3 one-model, overlap enabled + args_eagle_overlap = LlmArgs( + model="meta-llama/Meta-Llama-3.1-8B-Instruct", + speculative_config=spec_config, + disable_overlap_scheduler=False, + transforms=piecewise_disabled_transforms(), + ) + extra_overlap = get_extra_seq_len_for_kv_cache(args_eagle_overlap) + # Should be more than without overlap + assert extra_overlap > extra + + +def test_mtp_autodeploy_uses_eagle_one_model_capture(): + from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs + + model = "meta-llama/Meta-Llama-3.1-8B-Instruct" + args = LlmArgs( + model=model, + speculative_config=MTPDecodingConfig( + num_nextn_predict_layers=3, + mtp_eagle_one_model=True, + ), + transforms=piecewise_disabled_transforms(), + ) + + assert isinstance(args.speculative_config, MTPDecodingConfig) + assert args.model_factory == "eagle_one_model" + assert args.transforms["detect_hidden_states_for_capture"]["enabled"] is True + assert args.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] == {-1} + + +def test_detect_hidden_states_capture_last_layer_for_mtp_eagle_one_model(): + from tensorrt_llm._torch.auto_deploy.llm_args import LlmArgs + + config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") + config["args"].setdefault("transforms", {}).setdefault("compile_model", {})[ + "piecewise_enabled" + ] = False + + args = LlmArgs( + **config["args"], + speculative_config=MTPDecodingConfig( + num_nextn_predict_layers=3, + mtp_eagle_one_model=True, + speculative_model=config["args"]["model"], + ), + ) + + factory = args.create_factory() + assert isinstance(factory, EagleOneModelFactory) + + model = factory.target_factory.build_model("meta") + input_ids = torch.ones((1, 8), dtype=torch.int64) + position_ids = torch.arange(8, dtype=torch.int64).unsqueeze(0) + gm = torch_export_to_gm( + model, + args=(input_ids, position_ids), + ) + + transform = DetectHiddenStatesForCapture( + config=TransformConfig( + stage="pattern_matcher", + eagle3_layers_to_capture={-1}, + ) + ) + + original_residual_nodes = transform.collect_residual_add_nodes(gm) + assert original_residual_nodes + last_layer = max(original_residual_nodes) + last_layer_residual = original_residual_nodes[last_layer] + expected_arg_names = tuple( + arg.name if isinstance(arg, torch.fx.Node) else arg for arg in last_layer_residual.args + ) + + gm, info = transform._apply(gm, None, None, None) + + capture_nodes = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.auto_deploy.residual_add_for_capture.default + ] + + assert info.num_matches == 1 + assert len(capture_nodes) == 1 + capture_arg_names = tuple( + arg.name if isinstance(arg, torch.fx.Node) else arg for arg in capture_nodes[0].args + ) + assert capture_arg_names == expected_arg_names diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py new file mode 100644 index 000000000000..01456c01712e --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_ad_trtllm_serve.py @@ -0,0 +1,117 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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 multiprocessing as mp +import time + +import pytest +import yaml +from _model_test_utils import get_small_model_config # type: ignore +from click.testing import CliRunner +from openai import OpenAI + +from tensorrt_llm._utils import get_free_port +from tensorrt_llm.commands.serve import main as serve_main + + +def _run_serve_with_click(args): + runner = CliRunner() + # Blocks while server runs + result = runner.invoke(serve_main, args, catch_exceptions=False) + if result.exit_code != 0: + raise SystemExit(result.exit_code) + + +@pytest.mark.timeout(500) +def test_trtllm_serve_openai_chat_completion(tmp_path): + # Prepare small model config and extra options yaml + config = get_small_model_config("meta-llama/Meta-Llama-3.1-8B-Instruct") + extra_args = config["args"] + + # NOTE: trtllm attention backend fails on B200 (likely illegal memory access); use flashinfer. + extra_args["attn_backend"] = "flashinfer" + extra_options_path = tmp_path / "extra_llm_api_options.yaml" + with open(extra_options_path, "w") as f: + yaml.safe_dump(extra_args, f) + + host = "127.0.0.1" + port = get_free_port() + + # Use the same `model` string for server and client requests + model_id = extra_args["model"] + + args = [ + "serve", + f"{model_id}", + "--backend", + "_autodeploy", + "--host", + host, + "--port", + str(port), + "--extra_llm_api_options", + str(extra_options_path), + ] + + ctx = mp.get_context("spawn") + server = ctx.Process(target=_run_serve_with_click, args=(args,)) + server.start() + + try: + # Wait for server to be ready by polling /v1/models via OpenAI client + client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key="tensorrt_llm") + + start_time = time.time() + last_err = None + while time.time() - start_time < 300: + if not server.is_alive(): + raise RuntimeError("Server process exited prematurely") + try: + # Lightweight readiness probe + _ = client.models.list() + break + except Exception as e: # noqa: BLE001 + last_err = e + time.sleep(1) + else: + raise TimeoutError(f"Server did not become ready in time: {last_err}") + + # Send a small chat completion request + resp = client.chat.completions.create( + model=model_id, + messages=[ + {"role": "system", "content": "you are a helpful assistant"}, + {"role": "user", "content": "Say 'ok'"}, + ], + max_tokens=8, + ) + + # print response + print(f"{resp=}") + + assert hasattr(resp, "choices") and len(resp.choices) > 0 + first = resp.choices[0] + # new OpenAI client returns .message for chat completions + assert getattr(first, "message", None) is not None + # Content may be a string or a structured list depending on client version + _ = getattr(first.message, "content", None) + + finally: + # Terminate server and clean up + if server.is_alive(): + server.terminate() + server.join(timeout=20) + if server.is_alive(): + server.kill() + server.join(timeout=20) diff --git a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py index 490bcfd26570..10b802fe7173 100644 --- a/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py +++ b/tests/unittest/auto_deploy/singlegpu/smoke/test_disagg.py @@ -22,7 +22,10 @@ from tensorrt_llm import DisaggregatedParams, SamplingParams from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM +from tensorrt_llm.llmapi import Eagle3DecodingConfig +LLAMA_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct" +EAGLE3_MODEL_ID = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" DEEPSEEK_MODEL_ID = "deepseek-ai/DeepSeek-V3" DEEPSEEK_DISAGG_TRANSFORMS = { "insert_cached_attention": {"backend": "triton"}, @@ -97,6 +100,11 @@ def create_generation_params(context_output): return replace(params, request_type="generation_only") +def has_draft_tokens(output): + params = output.disaggregated_params + return params is not None and params.draft_tokens is not None and len(params.draft_tokens) > 0 + + def run_live_disagg_smoke( model_id, attn_backend, @@ -204,6 +212,16 @@ def run_live_batch_disagg_smoke(model_id, attn_backend, compile_backend, config_ GENERIC_DISAGG_SMOKE_CASES = [ + pytest.param(LLAMA_MODEL_ID, "trtllm", "torch-simple", {}, id="llama-trtllm-simple"), + pytest.param(LLAMA_MODEL_ID, "trtllm", "torch-cudagraph", {}, id="llama-trtllm-cudagraph"), + pytest.param(LLAMA_MODEL_ID, "flashinfer", "torch-simple", {}, id="llama-flashinfer-simple"), + pytest.param( + LLAMA_MODEL_ID, + "flashinfer", + "torch-cudagraph", + {}, + id="llama-flashinfer-cudagraph", + ), pytest.param( DEEPSEEK_MODEL_ID, "trtllm", @@ -243,3 +261,44 @@ def test_autodeploy_disaggregated_batch_smoke( ) run_live_batch_disagg_smoke(model_id, attn_backend, compile_backend, config_overrides) + + +def test_autodeploy_disaggregated_eagle3_smoke(): + target_model_config = get_small_model_config(LLAMA_MODEL_ID) + eagle3_model_config = get_small_model_config(EAGLE3_MODEL_ID) + target_model_kwargs = { + **target_model_config["args"]["model_kwargs"], + "num_hidden_layers": 3, + } + speculative_config = Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=eagle3_model_config["args"]["model"], + eagle3_one_model=True, + eagle3_layers_to_capture={0, 1, 2}, + ) + speculative_model_kwargs = { + **target_model_kwargs, + **eagle3_model_config["args"]["model_kwargs"], + "torch_dtype": "bfloat16", + } + + # This is intentionally a smoke test: small_model_config_disagg uses + # skip_loading_weights=True, so the meaningful assertions are that one-model + # Eagle builds with a reduced target/draft pair and carries draft-token + # metadata through the live disaggregated handoff. Force the draft dtype to + # match the BF16 Llama target because shared KV cache management requires + # target and draft KV resources to have the same dtype. Use three reduced + # target layers to match Llama Eagle3's default three-layer capture. + # Weighted acceptance and quality coverage belong in integration tests. + context_output, generation_output = run_live_disagg_smoke( + LLAMA_MODEL_ID, + "flashinfer", + "torch-simple", + common_config_overrides={ + "model_kwargs": target_model_kwargs, + "speculative_config": speculative_config, + "speculative_model_kwargs": speculative_model_kwargs, + }, + ) + assert has_draft_tokens(context_output) + assert has_draft_tokens(generation_output) From bceeaa5ab53b15a9c0001b40a7aaf706bbd86365 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 12:16:38 +0800 Subject: [PATCH 06/16] [None][test] Restore disaggregated Eagle3 handoff and logprobs tests on Qwen3-8B Re-add test_async_eagle3_full_model_handoff (test_ad_disagg.py) and test_disaggregated_logprobs_serving (test_disaggregated.py), both previously deleted as part of the Llama-3.1-8B retirement, using Qwen3-8B (+ Qwen3/qwen3_8b_eagle3 draft for the Eagle3 case) instead. test_async_eagle3_full_model_handoff now compares against a freshly computed aggregate run using the same speculative_config instead of hardcoded golden text/token IDs, since no such goldens exist yet for this model pairing. eagle3_layers_to_capture was scaled proportionally from the retired Llama config and is flagged as needing GPU validation. test_disaggregated_logprobs_serving drops the llama_model_root indirect fixture in favor of the simpler direct setup_model_symlink pattern already used elsewhere in this file, and adds a new disagg_config_ctxtp2_gentp2_qwen3_8b.yaml config. Both tests still need a real GPU run to confirm behavior before they can be trusted in CI. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- .../defs/disaggregated/test_ad_disagg.py | 79 +++++++ .../disagg_config_ctxtp2_gentp2_qwen3_8b.yaml | 18 ++ .../defs/disaggregated/test_disaggregated.py | 200 ++++++++++++++++++ 3 files changed, 297 insertions(+) create mode 100644 tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml diff --git a/tests/integration/defs/disaggregated/test_ad_disagg.py b/tests/integration/defs/disaggregated/test_ad_disagg.py index b4d49e845673..dae4927894fa 100644 --- a/tests/integration/defs/disaggregated/test_ad_disagg.py +++ b/tests/integration/defs/disaggregated/test_ad_disagg.py @@ -33,6 +33,7 @@ from tensorrt_llm import DisaggregatedParams, SamplingParams from tensorrt_llm._torch.auto_deploy import LLM as AutoDeployLLM from tensorrt_llm._utils import set_mpi_comm +from tensorrt_llm.llmapi import Eagle3DecodingConfig cloudpickle.register_pickle_by_value(sys.modules[__name__]) MPI.pickle.__init__( @@ -71,6 +72,8 @@ def skip_b300(): MODEL_PATHS = { "TinyLlama-1.1B-Chat-v1.0": "llama-models-v2/TinyLlama-1.1B-Chat-v1.0", "DeepSeek-V3-Lite": "DeepSeek-V3-Lite/bf16", + "Qwen3-8B-eagle3": "Qwen3/qwen3_8b_eagle3", + "Qwen3-8B": "Qwen3/Qwen3-8B", } @@ -635,6 +638,25 @@ def test_chunked_prefill_handoff(model): # --------------------------------------------------------------------------- +def qwen3_eagle3_config(): + return { + "speculative_config": Eagle3DecodingConfig( + max_draft_len=3, + speculative_model=model_path("Qwen3-8B-eagle3"), + eagle3_one_model=True, + # TODO: these capture layers were carried over proportionally from + # the retired Llama-3.1-8B config (1/32, 15/32, 28/32 through the + # stack) and have NOT been validated against Qwen3-8B's actual + # layer count on GPU. Re-derive and confirm before relying on this + # test's output. + eagle3_layers_to_capture={1, 17, 31}, + ), + # Force the Eagle3 draft to match the BF16 Qwen3-8B target. Shared KV + # cache management requires matching target and draft KV dtypes. + "speculative_model_kwargs": {"torch_dtype": "bfloat16"}, + } + + def get_ucx_tls() -> str: """Get UCX_TLS value based on GPU architecture. @@ -996,3 +1018,60 @@ def test_async_sharded_generation_handoff(): assert outputs["context"].token_ids == aggregate_output.token_ids[:1] assert outputs["generation"].text == aggregate_output.text assert outputs["generation"].token_ids == aggregate_output.token_ids + + +@skip_pre_hopper +@pytest.mark.threadleak(enabled=False) +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.skip_less_device(2) +@pytest.mark.timeout(900) +def test_async_eagle3_full_model_handoff(): + """Eagle3 one-model draft-token handoff, compared against an aggregate run. + + Unlike the retired Llama-3.1-8B version of this test, this compares + against a freshly-computed aggregate (non-disaggregated) generation using + the same speculative_config, instead of hardcoded golden text/token IDs. + That avoids needing pre-recorded goldens for the new model pairing, at the + cost of also exercising the aggregate Eagle3 one-model path as a + dependency. This still needs a real GPU run to confirm Qwen3-8B + + Qwen3/qwen3_8b_eagle3 actually produce matching, non-trivial draft-token + output under this config (see the eagle3_layers_to_capture TODO in + qwen3_eagle3_config). + """ + prompt = "What is the capital of Germany?" + sampling_params_kwargs = { + "max_tokens": 16, + "ignore_eos": True, + "top_k": 1, + "seed": AUTODEPLOY_DISAGG_SEED, + } + extra_config = qwen3_eagle3_config() + + aggregate_output = run_aggregate_generation( + "Qwen3-8B", + world_size=1, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + outputs = run_context_then_generation_handoff( + "Qwen3-8B", + worker_world_sizes=(1, 1), + generation_overlap=True, + prompt=prompt, + sampling_params_kwargs=sampling_params_kwargs, + extra_config=extra_config, + ) + context_params = outputs["context"].disaggregated_params + assert context_params is not None + assert context_params.request_type == "context_only" + assert len(outputs["context"].token_ids) == 1 + assert context_params.ctx_request_id is not None + assert context_params.first_gen_tokens is not None + assert has_handoff_transport_metadata(context_params) + assert outputs["generation"].token_ids + assert has_draft_tokens(outputs["context"]) + assert has_draft_tokens(outputs["generation"]) + assert outputs["context"].token_ids == aggregate_output.token_ids[:1] + assert outputs["generation"].text == aggregate_output.text + assert outputs["generation"].token_ids == aggregate_output.token_ids diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml new file mode 100644 index 000000000000..78b3614a650e --- /dev/null +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml @@ -0,0 +1,18 @@ +hostname: localhost +model: Qwen3/Qwen3-8B +free_gpu_memory_fraction: 0.25 +backend: pytorch +disable_overlap_scheduler: true +gather_generation_logits: true +context_servers: + num_instances: 1 + tensor_parallel_size: 2 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT +generation_servers: + num_instances: 1 + tensor_parallel_size: 2 + pipeline_parallel_size: 1 + cache_transceiver_config: + backend: DEFAULT diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index f521612400fa..d843b99c8ddd 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -28,6 +28,7 @@ from typing import Any, Optional import aiohttp +import numpy as np import pytest import yaml from defs.common import get_free_port_in_ci as get_free_port @@ -303,6 +304,8 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_multi_orchestrator.yaml", "4_ranks": f"{test_configs_root}/disagg_config_ctxtp2_gentp1.yaml", + "qwen3_8b": + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml", "cuda_graph": f"{test_configs_root}/disagg_config_cuda_graph_padding.yaml", "mixed": @@ -4045,6 +4048,203 @@ def run_disaggregated_cancel_test(example_dir, shutil.rmtree(work_dir, ignore_errors=True) +@pytest.mark.skip_less_device(4) +def test_disaggregated_logprobs_serving(disaggregated_test_root, + disaggregated_example_root, llm_venv): + """Test logprobs via OpenAI API in disaggregated serving with multi-GPU TP. + + Covers the RCCA scenario (NVBug 5926823): disaggregated + streaming + logprobs, + where the context worker returns prefill result (request_type=generation_only) + to the generation worker. Ensures LogProbStorage flows correctly across the + context/gen boundary without AttributeError on cum_log_probs. + """ + + async def iter_sse_chunks(resp): + """Yield parsed JSON chunks from an OpenAI SSE stream.""" + async for line in resp.content: + decoded = line.decode("utf-8").strip() + if not decoded.startswith("data: "): + continue + data_str = decoded[len("data: "):] + if data_str == "[DONE]": + break + try: + yield json.loads(data_str) + except json.JSONDecodeError: + continue + + async def collect_streaming_logprobs(resp, api_type): + """Parse SSE stream and return (tokens, logprobs) lists.""" + tokens, logprobs = [], [] + async for chunk in iter_sse_chunks(resp): + choices = chunk.get("choices", []) + if not choices: + continue + lp_data = choices[0].get("logprobs") + if not lp_data: + continue + if api_type == "completions": + tokens.extend(lp_data.get("tokens", [])) + logprobs.extend(lp_data.get("token_logprobs", [])) + else: + for item in lp_data.get("content", []): + tokens.append(item.get("token")) + logprobs.append(item.get("logprob")) + return tokens, logprobs + + def extract_logprobs(result, api_type): + """Extract (tokens, logprobs) from non-streaming OpenAI response.""" + choices = result.get("choices", []) + assert len(choices) > 0, "Response should have choices" + if api_type == "completions": + lp_data = choices[0].get("logprobs") + assert lp_data is not None, "Response should contain logprobs" + tokens = lp_data.get("tokens", []) + logprobs = lp_data.get("token_logprobs", []) + assert len(tokens) == len(logprobs), ( + f"count mismatch: {len(logprobs)} logprobs " + f"for {len(tokens)} tokens") + return tokens, logprobs + lp_obj = choices[0].get("logprobs") + assert lp_obj is not None, "Response should contain logprobs" + content = lp_obj.get("content", []) + tokens = [item.get("token") for item in content] + logprobs = [item.get("logprob") for item in content] + return tokens, logprobs + + model_path = "Qwen3/Qwen3-8B" + model_dir = f"{llm_models_root()}/{model_path}" + setup_model_symlink(llm_venv, model_dir, model_path) + + config_file = get_test_config("qwen3_8b", disaggregated_example_root, + os.path.dirname(__file__)) + + env = llm_venv._new_env.copy() + env["UCX_TLS"] = get_ucx_tls() + ctx_workers, gen_workers, disagg_server, work_dir = [], [], None, None + config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ + setup_disagg_cluster(config_file, env=env, + model_name=model_dir, + cwd=llm_venv.get_working_directory(), + server_start_timeout=600) + + server_host = config.get("hostname", "localhost") + server_url = f"http://{server_host}:{server_port}" + model_name = model_path + max_tokens = 20 + timeout = aiohttp.ClientTimeout(total=120) + # Use emoji prompt to also stress-test multi-byte tokenizer handling + prompt = "I love coding šŸš€ and AI." + + async def check_logprobs(): + async with aiohttp.ClientSession() as session: + for api_type in ("completions", "chat"): + url = (f"{server_url}/v1/completions" + if api_type == "completions" else + f"{server_url}/v1/chat/completions") + + def make_payload(prompt, stream, _api_type=api_type): + base = { + "max_tokens": max_tokens, + "logprobs": 1 if _api_type == "completions" else True, + "stream": stream, + "temperature": 0 + } + if _api_type == "completions": + return {"model": model_name, "prompt": prompt, **base} + return { + "model": model_name, + "messages": [{ + "role": "user", + "content": prompt + }], + **base + } + + # 1) Streaming vs non-streaming consistency check + async with session.post(url, + json=make_payload(prompt, False), + timeout=timeout) as resp: + assert resp.status == 200, \ + f"[{api_type}] non-streaming: {await resp.text()}" + ns_tokens, ns_logprobs = extract_logprobs( + await resp.json(), api_type) + + async with session.post(url, + json=make_payload(prompt, True), + timeout=timeout) as resp: + assert resp.status == 200, \ + f"[{api_type}] streaming: {await resp.text()}" + st_tokens, st_logprobs = \ + await collect_streaming_logprobs(resp, api_type) + + assert ns_tokens == st_tokens, ( + f"[{api_type}] streaming vs non-streaming tokens mismatch") + assert len(ns_logprobs) == len(st_logprobs), ( + f"[{api_type}] logprobs length: " + f"{len(ns_logprobs)} vs {len(st_logprobs)}") + # Skip position 0: the first token logprob can diverge + # between streaming and non-streaming in disaggregated mode + # due to the context/generation handoff boundary. + comparable = 0 + for i, (n, s) in enumerate( + zip(ns_logprobs, st_logprobs, strict=True)): + if i == 0 or n is None or s is None: + continue + comparable += 1 + rtol, atol = (1e-3, 1e-4) if api_type == "chat" else (1e-4, + 1e-5) + assert np.isclose(n, s, rtol=rtol, atol=atol), \ + f"[{api_type}] logprob mismatch at {i}: {n} vs {s}" + assert comparable > 0, ( + f"[{api_type}] no comparable post-handoff logprobs found") + + # 2) Chat API with top_logprobs (requires gather_generation_logits) + if api_type == "chat": + top_lp_payload = { + "model": model_name, + "messages": [{ + "role": "user", + "content": prompt + }], + "max_tokens": max_tokens, + "logprobs": True, + "top_logprobs": 3, + "stream": False, + "temperature": 0, + } + async with session.post(f"{server_url}/v1/chat/completions", + json=top_lp_payload, + timeout=timeout) as resp: + assert resp.status == 200, ( + f"[chat/top_logprobs] {resp.status}: " + f"{await resp.text()}") + result = await resp.json() + lp_obj = result["choices"][0].get("logprobs") + assert lp_obj is not None, "top_logprobs response should have logprobs" + content = lp_obj.get("content", []) + assert len( + content) > 0, "top_logprobs content should be non-empty" + for item in content: + top_lps = item.get("top_logprobs") + assert top_lps is not None and len(top_lps) > 0, ( + f"top_logprobs should be non-empty when requested: {item}" + ) + for tl in top_lps: + assert "token" in tl and "logprob" in tl, ( + f"top_logprob entry missing token/logprob: {tl}" + ) + assert tl["logprob"] <= 0.0, ( + f"top_logprob {tl['logprob']} should be <= 0") + + try: + asyncio.run(check_logprobs()) + finally: + terminate(*ctx_workers, *gen_workers, disagg_server) + if work_dir: + shutil.rmtree(work_dir, ignore_errors=True) + + @pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], indirect=True) def test_disaggregated_cancel_large_context_requests(disaggregated_test_root, From 6634e322059afd8a2479037f2d73b3aed4bcb647 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 12:59:56 +0800 Subject: [PATCH 07/16] [TRTLLM-15124][test] move test_disaggregated_logprobs_serving to end of file Keep it as a pure addition instead of an in-place replacement of the retired Llama-3.1-8B version, so it doesn't shift the line position of test_disaggregated_cancel_large_context_requests and everything after it. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- .../defs/disaggregated/test_disaggregated.py | 326 +++++++++--------- 1 file changed, 163 insertions(+), 163 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index d843b99c8ddd..a22ffe6d81c9 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -4048,6 +4048,169 @@ def run_disaggregated_cancel_test(example_dir, shutil.rmtree(work_dir, ignore_errors=True) +@pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], + indirect=True) +def test_disaggregated_cancel_large_context_requests(disaggregated_test_root, + disaggregated_example_root, + llm_venv, + deepseek_v3_model_root): + """ + Test that the disaggregated server handles request cancellations gracefully. + + This test sends bursts of requests with large contexts and cancels them + during prefill to stress test resource cleanup. + """ + setup_model_symlink(llm_venv, deepseek_v3_model_root, + "DeepSeek-V3-Lite/bf16") + + run_disaggregated_cancel_test(disaggregated_example_root, + "cancel_stress_test", + env=llm_venv._new_env, + num_bursts=5, + requests_per_burst=32, + model_path=deepseek_v3_model_root, + cwd=llm_venv.get_working_directory()) + + +@pytest.mark.skip_less_device(8) +@skip_pre_blackwell +@pytest.mark.parametrize("model_path", ['DeepSeek-V3-0324-FP4']) +def test_disaggregated_cancel_large_context_requests_long( + disaggregated_test_root, disaggregated_example_root, llm_venv, + model_path): + """Test that disaggregated server handles request cancellations gracefully. + + This test sends bursts of requests with large contexts and cancels them + during prefill to stress test resource cleanup. + """ + model_dir = f"{llm_models_root()}/{model_path}" + setup_model_symlink(llm_venv, model_dir, model_path) + + run_disaggregated_cancel_test(disaggregated_example_root, + "cancel_stress_test_large", + env=llm_venv._new_env, + num_bursts=1000, + requests_per_burst=32, + model_path=model_dir, + cwd=llm_venv.get_working_directory()) + + +@pytest.mark.skip_less_device(8) +@skip_pre_blackwell +@pytest.mark.parametrize("model_path", + ['NVIDIA-Nemotron-3-Super-120B-A12B-FP8']) +def test_disaggregated_mamba_conc_greater_than_mbs(disaggregated_example_root, + llm_venv, model_path, + benchmark_root, + shared_gpt_path): + model_dir = f"{llm_models_root()}/{model_path}" + setup_model_symlink(llm_venv, model_dir, model_path) + + config_file = get_test_config("mamba_conc_greater_than_mbs", + disaggregated_example_root, + os.path.dirname(__file__)) + + env = llm_venv._new_env.copy() + env["UCX_TLS"] = get_ucx_tls() + e2el, ttft = run_disaggregated_benchmark( + disaggregated_example_root, + config_file, + benchmark_root, + model_dir, + shared_gpt_path, + env=env, + num_prompts=40, + max_concurrency=4, + random_input_len=1024, + random_output_len=1024, + skip_warmup=True, + model_path=model_dir, + cwd=llm_venv.get_working_directory()) + print(f"E2EL: {e2el} ms, TTFT: {ttft} ms") + + +@pytest.mark.parametrize( + "test_config", + [ + # Smoke run: 120 requests at 64 concurrency (matching the warmup + # count), ~3 min request phase on B200. Used as L0 post-merge gate. A + # healthy cluster scores ~1.0 accuracy (every profile validates real + # content), so 0.9 leaves margin for a rare flaky request while still + # catching a regression. + pytest.param(TestConfig( + model_path='Qwen3/Qwen3-32B-FP8', + test_desc='req120-conc64-qwen3_32b_fp8_mixed_stress', + request_count=120, + concurrency=64, + accuracy_threshold=0.9, + speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), + marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), + # Full stress run: 10k requests at 512 concurrency. + # Estimated wall-clock ~40 min (server startup + ~32 min request + # phase); 512 concurrency exercises more in-flight overlap. + pytest.param(TestConfig( + model_path='Qwen3/Qwen3-32B-FP8', + test_desc='req10k-conc512-qwen3_32b_fp8_mixed_stress', + request_count=10000, + concurrency=512, + accuracy_threshold=0.9, + speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), + marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), + ], + ids=lambda x: x.test_desc) +def test_disaggregated_mixed_stress_test(disaggregated_test_root, + disaggregated_example_root, llm_venv, + test_config): + model_path = test_config.model_path + test_desc = test_config.test_desc + model_dir = resolve_llm_model_path(model_path) + setup_model_symlink(llm_venv, model_dir, model_path) + + config_file = get_test_config(test_desc, disaggregated_example_root, + os.path.dirname(__file__)) + + if test_config.speculative_model_path is not None: + spec_model_dir = f"{llm_models_root()}/{test_config.speculative_model_path}" + setup_model_symlink(llm_venv, spec_model_dir, + test_config.speculative_model_path) + with open(config_file, 'r') as f: + patched_config = yaml.safe_load(f) + patched_sections = [] + # Check top-level speculative_config first (current YAML layout), then + # fall back to per-server blocks for older config shapes. + top_spec = patched_config.get('speculative_config') + if isinstance(top_spec, dict) and 'speculative_model' in top_spec: + top_spec['speculative_model'] = spec_model_dir + patched_sections.append('top-level') + else: + for section in ('context_servers', 'generation_servers'): + spec = patched_config.get(section, {}).get('speculative_config') + if spec is not None and 'speculative_model' in spec: + spec['speculative_model'] = spec_model_dir + patched_sections.append(section) + if not patched_sections: + raise AssertionError( + f"{test_desc} sets speculative_model_path, but no " + "speculative_config.speculative_model field was patched") + patched_path = os.path.join(llm_venv.get_working_directory(), + f"{test_desc}_patched.yaml") + with open(patched_path, 'w') as f: + yaml.safe_dump(patched_config, f) + config_file = patched_path + + run_disaggregated_mixed_stress( + example_dir=disaggregated_example_root, + config_file=config_file, + model_path=model_dir, + total_requests=test_config.request_count, + concurrency=test_config.concurrency, + accuracy_threshold=test_config.accuracy_threshold, + incomplete_threshold=test_config.incomplete_threshold, + server_start_timeout=600, + env=llm_venv._new_env, + cwd=llm_venv.get_working_directory()) + + @pytest.mark.skip_less_device(4) def test_disaggregated_logprobs_serving(disaggregated_test_root, disaggregated_example_root, llm_venv): @@ -4243,166 +4406,3 @@ def make_payload(prompt, stream, _api_type=api_type): terminate(*ctx_workers, *gen_workers, disagg_server) if work_dir: shutil.rmtree(work_dir, ignore_errors=True) - - -@pytest.mark.parametrize("deepseek_v3_model_root", ['DeepSeek-V3-Lite-bf16'], - indirect=True) -def test_disaggregated_cancel_large_context_requests(disaggregated_test_root, - disaggregated_example_root, - llm_venv, - deepseek_v3_model_root): - """ - Test that the disaggregated server handles request cancellations gracefully. - - This test sends bursts of requests with large contexts and cancels them - during prefill to stress test resource cleanup. - """ - setup_model_symlink(llm_venv, deepseek_v3_model_root, - "DeepSeek-V3-Lite/bf16") - - run_disaggregated_cancel_test(disaggregated_example_root, - "cancel_stress_test", - env=llm_venv._new_env, - num_bursts=5, - requests_per_burst=32, - model_path=deepseek_v3_model_root, - cwd=llm_venv.get_working_directory()) - - -@pytest.mark.skip_less_device(8) -@skip_pre_blackwell -@pytest.mark.parametrize("model_path", ['DeepSeek-V3-0324-FP4']) -def test_disaggregated_cancel_large_context_requests_long( - disaggregated_test_root, disaggregated_example_root, llm_venv, - model_path): - """Test that disaggregated server handles request cancellations gracefully. - - This test sends bursts of requests with large contexts and cancels them - during prefill to stress test resource cleanup. - """ - model_dir = f"{llm_models_root()}/{model_path}" - setup_model_symlink(llm_venv, model_dir, model_path) - - run_disaggregated_cancel_test(disaggregated_example_root, - "cancel_stress_test_large", - env=llm_venv._new_env, - num_bursts=1000, - requests_per_burst=32, - model_path=model_dir, - cwd=llm_venv.get_working_directory()) - - -@pytest.mark.skip_less_device(8) -@skip_pre_blackwell -@pytest.mark.parametrize("model_path", - ['NVIDIA-Nemotron-3-Super-120B-A12B-FP8']) -def test_disaggregated_mamba_conc_greater_than_mbs(disaggregated_example_root, - llm_venv, model_path, - benchmark_root, - shared_gpt_path): - model_dir = f"{llm_models_root()}/{model_path}" - setup_model_symlink(llm_venv, model_dir, model_path) - - config_file = get_test_config("mamba_conc_greater_than_mbs", - disaggregated_example_root, - os.path.dirname(__file__)) - - env = llm_venv._new_env.copy() - env["UCX_TLS"] = get_ucx_tls() - e2el, ttft = run_disaggregated_benchmark( - disaggregated_example_root, - config_file, - benchmark_root, - model_dir, - shared_gpt_path, - env=env, - num_prompts=40, - max_concurrency=4, - random_input_len=1024, - random_output_len=1024, - skip_warmup=True, - model_path=model_dir, - cwd=llm_venv.get_working_directory()) - print(f"E2EL: {e2el} ms, TTFT: {ttft} ms") - - -@pytest.mark.parametrize( - "test_config", - [ - # Smoke run: 120 requests at 64 concurrency (matching the warmup - # count), ~3 min request phase on B200. Used as L0 post-merge gate. A - # healthy cluster scores ~1.0 accuracy (every profile validates real - # content), so 0.9 leaves margin for a rare flaky request while still - # catching a regression. - pytest.param(TestConfig( - model_path='Qwen3/Qwen3-32B-FP8', - test_desc='req120-conc64-qwen3_32b_fp8_mixed_stress', - request_count=120, - concurrency=64, - accuracy_threshold=0.9, - speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), - marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), - # Full stress run: 10k requests at 512 concurrency. - # Estimated wall-clock ~40 min (server startup + ~32 min request - # phase); 512 concurrency exercises more in-flight overlap. - pytest.param(TestConfig( - model_path='Qwen3/Qwen3-32B-FP8', - test_desc='req10k-conc512-qwen3_32b_fp8_mixed_stress', - request_count=10000, - concurrency=512, - accuracy_threshold=0.9, - speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), - marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), - ], - ids=lambda x: x.test_desc) -def test_disaggregated_mixed_stress_test(disaggregated_test_root, - disaggregated_example_root, llm_venv, - test_config): - model_path = test_config.model_path - test_desc = test_config.test_desc - model_dir = resolve_llm_model_path(model_path) - setup_model_symlink(llm_venv, model_dir, model_path) - - config_file = get_test_config(test_desc, disaggregated_example_root, - os.path.dirname(__file__)) - - if test_config.speculative_model_path is not None: - spec_model_dir = f"{llm_models_root()}/{test_config.speculative_model_path}" - setup_model_symlink(llm_venv, spec_model_dir, - test_config.speculative_model_path) - with open(config_file, 'r') as f: - patched_config = yaml.safe_load(f) - patched_sections = [] - # Check top-level speculative_config first (current YAML layout), then - # fall back to per-server blocks for older config shapes. - top_spec = patched_config.get('speculative_config') - if isinstance(top_spec, dict) and 'speculative_model' in top_spec: - top_spec['speculative_model'] = spec_model_dir - patched_sections.append('top-level') - else: - for section in ('context_servers', 'generation_servers'): - spec = patched_config.get(section, {}).get('speculative_config') - if spec is not None and 'speculative_model' in spec: - spec['speculative_model'] = spec_model_dir - patched_sections.append(section) - if not patched_sections: - raise AssertionError( - f"{test_desc} sets speculative_model_path, but no " - "speculative_config.speculative_model field was patched") - patched_path = os.path.join(llm_venv.get_working_directory(), - f"{test_desc}_patched.yaml") - with open(patched_path, 'w') as f: - yaml.safe_dump(patched_config, f) - config_file = patched_path - - run_disaggregated_mixed_stress( - example_dir=disaggregated_example_root, - config_file=config_file, - model_path=model_dir, - total_requests=test_config.request_count, - concurrency=test_config.concurrency, - accuracy_threshold=test_config.accuracy_threshold, - incomplete_threshold=test_config.incomplete_threshold, - server_start_timeout=600, - env=llm_venv._new_env, - cwd=llm_venv.get_working_directory()) From 1f0f7473be82400f90d2249779c9c597a40f0293 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:14:53 +0800 Subject: [PATCH 08/16] [TRTLLM-15124][test] replace Llama-3.1-8B with Qwen3-8B in test_disaggregated_logprobs_serving Rebuild the test from upstream's current version and swap only the model-specific pieces, so the diff stays minimal and no other test in the file shifts position. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- .../defs/disaggregated/test_disaggregated.py | 282 +++++++++--------- 1 file changed, 141 insertions(+), 141 deletions(-) diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index a22ffe6d81c9..632806005168 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -304,8 +304,6 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_multi_orchestrator.yaml", "4_ranks": f"{test_configs_root}/disagg_config_ctxtp2_gentp1.yaml", - "qwen3_8b": - f"{test_configs_root}/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml", "cuda_graph": f"{test_configs_root}/disagg_config_cuda_graph_padding.yaml", "mixed": @@ -422,6 +420,8 @@ def get_test_config(test_desc, example_dir, test_root): f"{test_configs_root}/disagg_config_cancel_stress_test.yaml", "cancel_stress_test_large": f"{test_configs_root}/disagg_config_cancel_stress_test_large.yaml", + "qwen3_8b": + f"{test_configs_root}/disagg_config_ctxtp2_gentp2_qwen3_8b.yaml", "mamba_conc_greater_than_mbs": f"{test_configs_root}/disagg_config_mamba_conc_greater_than_mbs.yaml", "mamba_bs1_concurrency2": @@ -4072,145 +4072,6 @@ def test_disaggregated_cancel_large_context_requests(disaggregated_test_root, cwd=llm_venv.get_working_directory()) -@pytest.mark.skip_less_device(8) -@skip_pre_blackwell -@pytest.mark.parametrize("model_path", ['DeepSeek-V3-0324-FP4']) -def test_disaggregated_cancel_large_context_requests_long( - disaggregated_test_root, disaggregated_example_root, llm_venv, - model_path): - """Test that disaggregated server handles request cancellations gracefully. - - This test sends bursts of requests with large contexts and cancels them - during prefill to stress test resource cleanup. - """ - model_dir = f"{llm_models_root()}/{model_path}" - setup_model_symlink(llm_venv, model_dir, model_path) - - run_disaggregated_cancel_test(disaggregated_example_root, - "cancel_stress_test_large", - env=llm_venv._new_env, - num_bursts=1000, - requests_per_burst=32, - model_path=model_dir, - cwd=llm_venv.get_working_directory()) - - -@pytest.mark.skip_less_device(8) -@skip_pre_blackwell -@pytest.mark.parametrize("model_path", - ['NVIDIA-Nemotron-3-Super-120B-A12B-FP8']) -def test_disaggregated_mamba_conc_greater_than_mbs(disaggregated_example_root, - llm_venv, model_path, - benchmark_root, - shared_gpt_path): - model_dir = f"{llm_models_root()}/{model_path}" - setup_model_symlink(llm_venv, model_dir, model_path) - - config_file = get_test_config("mamba_conc_greater_than_mbs", - disaggregated_example_root, - os.path.dirname(__file__)) - - env = llm_venv._new_env.copy() - env["UCX_TLS"] = get_ucx_tls() - e2el, ttft = run_disaggregated_benchmark( - disaggregated_example_root, - config_file, - benchmark_root, - model_dir, - shared_gpt_path, - env=env, - num_prompts=40, - max_concurrency=4, - random_input_len=1024, - random_output_len=1024, - skip_warmup=True, - model_path=model_dir, - cwd=llm_venv.get_working_directory()) - print(f"E2EL: {e2el} ms, TTFT: {ttft} ms") - - -@pytest.mark.parametrize( - "test_config", - [ - # Smoke run: 120 requests at 64 concurrency (matching the warmup - # count), ~3 min request phase on B200. Used as L0 post-merge gate. A - # healthy cluster scores ~1.0 accuracy (every profile validates real - # content), so 0.9 leaves margin for a rare flaky request while still - # catching a regression. - pytest.param(TestConfig( - model_path='Qwen3/Qwen3-32B-FP8', - test_desc='req120-conc64-qwen3_32b_fp8_mixed_stress', - request_count=120, - concurrency=64, - accuracy_threshold=0.9, - speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), - marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), - # Full stress run: 10k requests at 512 concurrency. - # Estimated wall-clock ~40 min (server startup + ~32 min request - # phase); 512 concurrency exercises more in-flight overlap. - pytest.param(TestConfig( - model_path='Qwen3/Qwen3-32B-FP8', - test_desc='req10k-conc512-qwen3_32b_fp8_mixed_stress', - request_count=10000, - concurrency=512, - accuracy_threshold=0.9, - speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), - marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), - ], - ids=lambda x: x.test_desc) -def test_disaggregated_mixed_stress_test(disaggregated_test_root, - disaggregated_example_root, llm_venv, - test_config): - model_path = test_config.model_path - test_desc = test_config.test_desc - model_dir = resolve_llm_model_path(model_path) - setup_model_symlink(llm_venv, model_dir, model_path) - - config_file = get_test_config(test_desc, disaggregated_example_root, - os.path.dirname(__file__)) - - if test_config.speculative_model_path is not None: - spec_model_dir = f"{llm_models_root()}/{test_config.speculative_model_path}" - setup_model_symlink(llm_venv, spec_model_dir, - test_config.speculative_model_path) - with open(config_file, 'r') as f: - patched_config = yaml.safe_load(f) - patched_sections = [] - # Check top-level speculative_config first (current YAML layout), then - # fall back to per-server blocks for older config shapes. - top_spec = patched_config.get('speculative_config') - if isinstance(top_spec, dict) and 'speculative_model' in top_spec: - top_spec['speculative_model'] = spec_model_dir - patched_sections.append('top-level') - else: - for section in ('context_servers', 'generation_servers'): - spec = patched_config.get(section, {}).get('speculative_config') - if spec is not None and 'speculative_model' in spec: - spec['speculative_model'] = spec_model_dir - patched_sections.append(section) - if not patched_sections: - raise AssertionError( - f"{test_desc} sets speculative_model_path, but no " - "speculative_config.speculative_model field was patched") - patched_path = os.path.join(llm_venv.get_working_directory(), - f"{test_desc}_patched.yaml") - with open(patched_path, 'w') as f: - yaml.safe_dump(patched_config, f) - config_file = patched_path - - run_disaggregated_mixed_stress( - example_dir=disaggregated_example_root, - config_file=config_file, - model_path=model_dir, - total_requests=test_config.request_count, - concurrency=test_config.concurrency, - accuracy_threshold=test_config.accuracy_threshold, - incomplete_threshold=test_config.incomplete_threshold, - server_start_timeout=600, - env=llm_venv._new_env, - cwd=llm_venv.get_working_directory()) - - @pytest.mark.skip_less_device(4) def test_disaggregated_logprobs_serving(disaggregated_test_root, disaggregated_example_root, llm_venv): @@ -4406,3 +4267,142 @@ def make_payload(prompt, stream, _api_type=api_type): terminate(*ctx_workers, *gen_workers, disagg_server) if work_dir: shutil.rmtree(work_dir, ignore_errors=True) + + +@pytest.mark.skip_less_device(8) +@skip_pre_blackwell +@pytest.mark.parametrize("model_path", ['DeepSeek-V3-0324-FP4']) +def test_disaggregated_cancel_large_context_requests_long( + disaggregated_test_root, disaggregated_example_root, llm_venv, + model_path): + """Test that disaggregated server handles request cancellations gracefully. + + This test sends bursts of requests with large contexts and cancels them + during prefill to stress test resource cleanup. + """ + model_dir = f"{llm_models_root()}/{model_path}" + setup_model_symlink(llm_venv, model_dir, model_path) + + run_disaggregated_cancel_test(disaggregated_example_root, + "cancel_stress_test_large", + env=llm_venv._new_env, + num_bursts=1000, + requests_per_burst=32, + model_path=model_dir, + cwd=llm_venv.get_working_directory()) + + +@pytest.mark.skip_less_device(8) +@skip_pre_blackwell +@pytest.mark.parametrize("model_path", + ['NVIDIA-Nemotron-3-Super-120B-A12B-FP8']) +def test_disaggregated_mamba_conc_greater_than_mbs(disaggregated_example_root, + llm_venv, model_path, + benchmark_root, + shared_gpt_path): + model_dir = f"{llm_models_root()}/{model_path}" + setup_model_symlink(llm_venv, model_dir, model_path) + + config_file = get_test_config("mamba_conc_greater_than_mbs", + disaggregated_example_root, + os.path.dirname(__file__)) + + env = llm_venv._new_env.copy() + env["UCX_TLS"] = get_ucx_tls() + e2el, ttft = run_disaggregated_benchmark( + disaggregated_example_root, + config_file, + benchmark_root, + model_dir, + shared_gpt_path, + env=env, + num_prompts=40, + max_concurrency=4, + random_input_len=1024, + random_output_len=1024, + skip_warmup=True, + model_path=model_dir, + cwd=llm_venv.get_working_directory()) + print(f"E2EL: {e2el} ms, TTFT: {ttft} ms") + + +@pytest.mark.parametrize( + "test_config", + [ + # Smoke run: 120 requests at 64 concurrency (matching the warmup + # count), ~3 min request phase on B200. Used as L0 post-merge gate. A + # healthy cluster scores ~1.0 accuracy (every profile validates real + # content), so 0.9 leaves margin for a rare flaky request while still + # catching a regression. + pytest.param(TestConfig( + model_path='Qwen3/Qwen3-32B-FP8', + test_desc='req120-conc64-qwen3_32b_fp8_mixed_stress', + request_count=120, + concurrency=64, + accuracy_threshold=0.9, + speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), + marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), + # Full stress run: 10k requests at 512 concurrency. + # Estimated wall-clock ~40 min (server startup + ~32 min request + # phase); 512 concurrency exercises more in-flight overlap. + pytest.param(TestConfig( + model_path='Qwen3/Qwen3-32B-FP8', + test_desc='req10k-conc512-qwen3_32b_fp8_mixed_stress', + request_count=10000, + concurrency=512, + accuracy_threshold=0.9, + speculative_model_path='Zhi-Create-Qwen3-32B-Eagle3'), + marks=(pytest.mark.skip_less_device(8), skip_pre_hopper)), + ], + ids=lambda x: x.test_desc) +def test_disaggregated_mixed_stress_test(disaggregated_test_root, + disaggregated_example_root, llm_venv, + test_config): + model_path = test_config.model_path + test_desc = test_config.test_desc + model_dir = resolve_llm_model_path(model_path) + setup_model_symlink(llm_venv, model_dir, model_path) + + config_file = get_test_config(test_desc, disaggregated_example_root, + os.path.dirname(__file__)) + + if test_config.speculative_model_path is not None: + spec_model_dir = f"{llm_models_root()}/{test_config.speculative_model_path}" + setup_model_symlink(llm_venv, spec_model_dir, + test_config.speculative_model_path) + with open(config_file, 'r') as f: + patched_config = yaml.safe_load(f) + patched_sections = [] + # Check top-level speculative_config first (current YAML layout), then + # fall back to per-server blocks for older config shapes. + top_spec = patched_config.get('speculative_config') + if isinstance(top_spec, dict) and 'speculative_model' in top_spec: + top_spec['speculative_model'] = spec_model_dir + patched_sections.append('top-level') + else: + for section in ('context_servers', 'generation_servers'): + spec = patched_config.get(section, {}).get('speculative_config') + if spec is not None and 'speculative_model' in spec: + spec['speculative_model'] = spec_model_dir + patched_sections.append(section) + if not patched_sections: + raise AssertionError( + f"{test_desc} sets speculative_model_path, but no " + "speculative_config.speculative_model field was patched") + patched_path = os.path.join(llm_venv.get_working_directory(), + f"{test_desc}_patched.yaml") + with open(patched_path, 'w') as f: + yaml.safe_dump(patched_config, f) + config_file = patched_path + + run_disaggregated_mixed_stress( + example_dir=disaggregated_example_root, + config_file=config_file, + model_path=model_dir, + total_requests=test_config.request_count, + concurrency=test_config.concurrency, + accuracy_threshold=test_config.accuracy_threshold, + incomplete_threshold=test_config.incomplete_threshold, + server_start_timeout=600, + env=llm_venv._new_env, + cwd=llm_venv.get_working_directory()) From b322a6dac54b8f560c4161440e09f1fa22681afd Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:27:57 +0800 Subject: [PATCH 09/16] [TRTLLM-15124][test] replace Llama-3.1-8B with Qwen3-8B in test_disaggregated_spec_dec_batch_slot_limit Restore this Eagle3 disaggregated batch-slot-release test on Qwen3-8B instead of the retired Llama-3.1-8B + EAGLE3-LLaMA3.1-Instruct-8B pairing, and add back its l0_h100 and QA test-list entries under the new parametrize ID. Also restores the renamed test_disaggregated_logprobs_serving entry in llm_function_core.txt. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- .../test_disaggregated_single_gpu.py | 112 ++++++++++++++++++ .../test_lists/qa/llm_function_core.txt | 2 + .../test_lists/test-db/l0_h100.yml | 1 + 3 files changed, 115 insertions(+) diff --git a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py index 59805c24e67c..5602aad62e58 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py +++ b/tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py @@ -14,6 +14,7 @@ from tensorrt_llm._utils import set_mpi_comm from tensorrt_llm.llmapi import (CacheTransceiverConfig, CudaGraphConfig, KvCacheConfig, MpiCommSession) +from tensorrt_llm.llmapi.llm_args import Eagle3DecodingConfig def get_ucx_tls(): @@ -48,7 +49,9 @@ def get_ucx_tls(): MODEL_PATHS = { "DeepSeek-V3-Lite-fp8": "DeepSeek-V3-Lite/fp8", "TinyLlama-1.1B-Chat-v1.0": "llama-models-v2/TinyLlama-1.1B-Chat-v1.0", + "Qwen3-8B-eagle3": "Qwen3/qwen3_8b_eagle3", "Qwen3-8B-FP8": "Qwen3/Qwen3-8B-FP8", + "Qwen3-8B": "Qwen3/Qwen3-8B", } @@ -507,6 +510,115 @@ def test_disaggregated_llama_context_capacity(model, enable_cuda_graph, print("All workers terminated.") +@skip_pre_hopper +@pytest.mark.parametrize("model", ["Qwen3-8B"]) +@pytest.mark.parametrize("spec_dec_model_path", ["Qwen3-8B-eagle3"]) +@pytest.mark.parametrize("generation_overlap", [False]) +def test_disaggregated_spec_dec_batch_slot_limit(model, spec_dec_model_path, + generation_overlap): + # Test whether the batch slots are properly released when using speculative decoding + # with disaggregated serving. + spec_dec_config = Eagle3DecodingConfig( + speculative_model=model_path(spec_dec_model_path), max_draft_len=3) + + worker_pytorch_configs = [] + + # Context worker + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=True, + speculative_config=spec_dec_config, + max_batch_size=1)) + + # Generation worker + worker_pytorch_configs.append( + dict(disable_overlap_scheduler=not generation_overlap, + speculative_config=spec_dec_config, + max_batch_size=1)) + + kv_cache_configs = [ + KvCacheConfig(max_tokens=128, + enable_block_reuse=False, + free_gpu_memory_fraction=0.4) for _ in range(2) + ] + cache_transceiver_configs = [ + CacheTransceiverConfig(backend="DEFAULT") for _ in range(2) + ] + model_names = [model_path(model) for _ in range(2)] + ranks = [0, 1] + worker_args = list( + zip(kv_cache_configs, cache_transceiver_configs, worker_pytorch_configs, + model_names, ranks)) + + port_name = mpi_publish_name() + + prompt = "What is the capital of Germany?" + mpi_info = MPI.Info.Create() + mpi_info.Set("oversubscribe", "true") + with MPIPoolExecutor(max_workers=2, + env={ + "UCX_TLS": get_ucx_tls(), + "UCX_MM_ERROR_HANDLING": "y", + "OMPI_MCA_rmaps_base_oversubscribe": "1" + }, + mpi_info=mpi_info) as executor: + futures = [] + try: + for worker_arg in worker_args: + future = executor.submit(worker_entry_point, *worker_arg) + futures.append(future) + except Exception as e: + print(f"Error in worker {worker_arg}: {e}") + raise e + + intercomm = None + try: + print("Launched all the workers.") + intercomm = mpi_initialize_intercomm(port_name) + + for _ in range(2): + intercomm.recv(tag=MPI_READY) + print("Received ready signal.") + max_tokens = 25 + + requests = [] + for _ in range(10): + requests.append( + (prompt, SamplingParams(max_tokens=1, ignore_eos=True), + DisaggregatedParams(request_type="context_only"))) + + intercomm.send(requests, dest=0, tag=MPI_REQUEST) + + for _ in range(len(requests)): + output = intercomm.recv(source=0, tag=MPI_RESULT) + assert output[0].disaggregated_params is not None + assert output[ + 0].disaggregated_params.request_type == "context_only" + assert len(output[0].token_ids) == 1 + + generation_request_disagg_params = output[ + 0].disaggregated_params + generation_request_disagg_params.request_type = "generation_only" + requests = [] + requests.append((prompt, + SamplingParams(max_tokens=max_tokens, + ignore_eos=True), + generation_request_disagg_params)) + + intercomm.send(requests, dest=1, tag=MPI_REQUEST) + output = intercomm.recv(source=1, tag=MPI_RESULT) + + except MPI.Exception as e: + print(f"MPI Error") + raise e + finally: + mpi_send_termination_request(intercomm) + + # Wait for all futures to complete + for future in futures: + future.result() + print("All workers terminated.") + + @pytest.mark.parametrize("model", ["TinyLlama-1.1B-Chat-v1.0"]) @pytest.mark.parametrize("generation_overlap", [False, True]) def test_disaggregated_logprobs(model, generation_overlap): diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 8dd23f691909..2f857d5c54a3 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -735,6 +735,7 @@ disaggregated/test_disaggregated.py::test_disaggregated_genbs1[TinyLlama-1.1B-Ch disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b] disaggregated/test_disaggregated.py::test_disaggregated_kv_cache_time_output[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0] +disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving disaggregated/test_disaggregated.py::test_disaggregated_mamba_bs1_concurrency2 disaggregated/test_disaggregated.py::test_disaggregated_mamba_conc_greater_than_mbs[NVIDIA-Nemotron-3-Super-120B-A12B-FP8] disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0] @@ -774,6 +775,7 @@ disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[ disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[False-True-Qwen3-8B-FP8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-False-Qwen3-8B-FP8] disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] +disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-Qwen3-8B-eagle3-Qwen3-8B] disaggregated/test_workers.py::test_workers_conditional_disaggregation[TinyLlama-1.1B-Chat-v1.0] disaggregated/test_workers.py::test_workers_conditional_disaggregation_deepseek_v3_lite_bf16[DeepSeek-V3-Lite-bf16] disaggregated/test_workers.py::test_workers_conversation_router[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 619c89deddf3..d86a9904314b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -173,6 +173,7 @@ l0_h100: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-False-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] + - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-Qwen3-8B-eagle3-Qwen3-8B] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_cancel_gen_requests[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[False-TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[True-TinyLlama-1.1B-Chat-v1.0] From 0e170592b58727b9f4c00e5c0238aa0d5d34490f Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:38:48 +0800 Subject: [PATCH 10/16] [TRTLLM-15124][test] replace Llama-3.1-8B with Qwen3.5-4B in multi-node disagg completion tests Restore test_openai_disagg_multi_nodes_completion, test_openai_disagg_multi_nodes_completion_service_discovery, and their underlying _test_disagg_serving_multi_nodes*.py helpers on Qwen3.5-4B instead of the retired Llama-3.1-8B, adding skip_pre_hopper since Qwen3.5-4B requires it. Also restores the three QA test-list entries in llm_function_multinode.txt. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/defs/test_e2e.py | 35 +++ .../test_lists/qa/llm_function_multinode.txt | 3 + .../apps/_test_disagg_serving_multi_nodes.py | 213 +++++++++++++++++ ...g_serving_multi_nodes_service_discovery.py | 223 ++++++++++++++++++ 4 files changed, 474 insertions(+) create mode 100644 tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py create mode 100644 tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 7d32cfed017f..7388f61e1c09 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -652,6 +652,41 @@ def test_trtllm_multimodal_benchmark_serving(llm_root, llm_venv): ]) +@skip_pre_hopper +@pytest.mark.skip_less_device(4) +@pytest.mark.skip_less_device_memory(40000) +@pytest.mark.parametrize("service_discovery", ["etcd"]) +def test_openai_disagg_multi_nodes_completion_service_discovery( + llm_root, llm_venv, service_discovery): + test_root = unittest_path() / "llmapi" / "apps" + llm_venv.run_cmd([ + "-m", + "pytest", + str(test_root / + f"_test_disagg_serving_multi_nodes_service_discovery.py::test_completion[{service_discovery}]" + ), + ]) + + +@skip_pre_hopper +@pytest.mark.skip_less_device(4) +@pytest.mark.skip_less_device_memory(40000) +@pytest.mark.parametrize("gen_config", + ["gen_tp2pp1", "gen_tp1pp2", "gen_tp1pp1"]) +@pytest.mark.parametrize("ctx_config", + ["ctx_tp2pp1", "ctx_tp1pp2", "ctx_tp1pp1"]) +def test_openai_disagg_multi_nodes_completion(llm_root, llm_venv, ctx_config, + gen_config): + test_root = unittest_path() / "llmapi" / "apps" + llm_venv.run_cmd([ + "-m", + "pytest", + str(test_root / + f"_test_disagg_serving_multi_nodes.py::test_completion[{ctx_config}-{gen_config}]" + ), + ]) + + ### PyTorch examples diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 836c5bd48a0e..c460bdda46dc 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -5,6 +5,9 @@ accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_nvfp4_adp16_cuteds test_e2e.py::test_multi_nodes_eval[DeepSeek-R1/DeepSeek-R1-0528-FP4-tp16-mmlu] test_e2e.py::test_multi_nodes_eval[Kimi-K2-Thinking-NVFP4-tp16-mmlu] test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] +test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] +test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] +test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[baseline] accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[dspark] accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[reuse] diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py new file mode 100644 index 000000000000..2324e5a5aece --- /dev/null +++ b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py @@ -0,0 +1,213 @@ +import os +import platform +import socket +import time + +import openai +import pytest +from test_common.http_utils import (wait_for_endpoint_down, + wait_for_endpoint_ready) + +from tensorrt_llm._utils import get_sm_version + +from ..test_llm import get_model_path +from .openai_server import RemoteDisaggOpenAIServer, RemoteOpenAIServer +from .utils import (expand_slurm_nodelist, wait_for_endpoint_down, + wait_for_endpoint_ready) + +RANK = int(os.environ.get("SLURM_PROCID", 0)) +NODE_RANK = int(os.environ.get("SLURM_NODEID", 0)) +NODE_LIST = expand_slurm_nodelist(os.environ.get("SLURM_NODELIST", "")) +SLURM_NTASKS_PER_NODE = int(os.environ.get("SLURM_NTASKS_PER_NODE", 1)) + +pytestmark = pytest.mark.threadleak(enabled=False) + +# This test assumes that there are >2 nodes, we run ctx/disagg-server/client on the first node, +# and run gen the second node. +# This is a multi-node test, and will not be scheduled to the same node running other tests +# using fixed ports should be safe. +CTX_SERVER_PORT = 8001 +GEN_SERVER_PORT = 8002 +DISAGG_SERVER_PORT = 8000 + + +# Exclude the current node from the node list, then return other nodes by idx +def get_the_other_host(idx=0): + assert len(NODE_LIST) >= 2 + node_list = NODE_LIST.copy() + curr_host = socket.gethostname() + if curr_host in NODE_LIST: + # gethostname returns the exact node name in node list + node_list.remove(curr_host) + else: + # gethostname returns the full domain + curr_host = curr_host.split('.')[0] + assert curr_host in node_list + node_list.remove(curr_host) + return node_list[idx] + + +def is_ctx_node(): + return NODE_RANK == 0 + + +def is_gen_node(): + return NODE_RANK == 1 + + +def is_disagg_node(): + return NODE_RANK == 0 + + +# The test is run on multinodes but only the first node's output is used for assertion +def is_pytest_node(): + return NODE_RANK == 0 + + +def env(): + # Remove MPI related environment variables to isolate the ctx/gen processes + # so that they will not be in the same MPI communicator, otherwise the rank and world_size may mismatch + e = { + k: v + for k, v in os.environ.items() + if not ('PMI_' in k or 'OMPI_' in k or 'PMIX_' in k or 'SLURM_' in k) + and k not in ["UCX_TLS", "UCX_NET_DEVICES"] # avoid UCX failure on oci + } + # Some GB300 machines have NICs that misbehave with UCX's default transport + # auto-selection, so UCX_TLS must be set explicitly there. Identify GB300 via + # sm_103 (Blackwell Ultra) + aarch64 (Grace) -- this excludes HGX B300 + # (also sm_103 but x86_64) and avoids depending on the GPU device name string. + if get_sm_version() == 103 and platform.machine().lower() == "aarch64": + e["UCX_TLS"] = "cuda_copy,cuda_ipc,sm,self,tcp" + return e + + +@pytest.fixture(scope="module") +def model_name(): + return "Qwen3.5-4B" + + +@pytest.fixture(scope="module", params=['pytorch'], ids=["pytorch"]) +def backend(request): + return request.param + + +@pytest.fixture( + scope="module", + params=[(1, 1), (2, 1), (1, 2)], + ids=lambda tp_pp_size: f'ctx_tp{tp_pp_size[0]}pp{tp_pp_size[1]}') +def ctx_tp_pp_size(request): + return request.param + + +@pytest.fixture( + scope="module", + params=[(1, 1), (2, 1), (1, 2)], + ids=lambda tp_pp_size: f'gen_tp{tp_pp_size[0]}pp{tp_pp_size[1]}') +def gen_tp_pp_size(request): + return request.param + + +@pytest.fixture(scope="module") +def worker(model_name: str, ctx_tp_pp_size: tuple, gen_tp_pp_size: tuple): + extra_config = { + "cache_transceiver_config": { + "backend": "DEFAULT" + }, + "kv_cache_config": { + "free_gpu_memory_fraction": 0.5, + "enable_block_reuse": False, + }, + "disable_overlap_scheduler": True, + } + if is_ctx_node(): + print(f"starting ctx_server for rank {RANK} node rank {NODE_RANK}") + model_path = get_model_path(model_name) + tp_size, pp_size = ctx_tp_pp_size + args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] + with RemoteOpenAIServer(model_path, + port=CTX_SERVER_PORT, + cli_args=args, + host="0.0.0.0", + env=env(), + llmapi_launch=False, + rank=RANK % SLURM_NTASKS_PER_NODE, + extra_config=extra_config) as server: + yield server + elif is_gen_node(): + print(f"starting gen_server for rank {RANK} node rank {NODE_RANK}") + model_path = get_model_path(model_name) + tp_size, pp_size = gen_tp_pp_size + args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] + with RemoteOpenAIServer(model_path, + port=GEN_SERVER_PORT, + cli_args=args, + host="0.0.0.0", + env=env(), + llmapi_launch=False, + rank=RANK % SLURM_NTASKS_PER_NODE, + extra_config=extra_config) as server: + yield server + else: + yield None + + +@pytest.fixture(scope="module") +def disagg_server(worker: RemoteOpenAIServer): + if is_disagg_node(): + print(f"starting disagg_server for rank {RANK} node rank {NODE_RANK}") + ctx_url = f"localhost:8001" # Use localhost since the ctx server is on the same node + # TODO: Hopefully the NODE_LIST is ordered by NODE_RANK, this test is only expected to run with 2 nodes now + # We need to test with 4 nodes or more in the future, which should be easier with service discovery + gen_url = f"{get_the_other_host(0)}:8002" + with RemoteDisaggOpenAIServer(ctx_servers=[ctx_url], + gen_servers=[gen_url], + port=DISAGG_SERVER_PORT, + llmapi_launch=False, + env=env()) as server: + yield server + else: + print(f"skipping disagg_server for rank {RANK} node rank {NODE_RANK}") + yield None + + +@pytest.fixture(scope="module") +def client(disagg_server: RemoteDisaggOpenAIServer): + if is_pytest_node(): + return disagg_server.get_client() + else: + print(f"skipping client for rank {RANK} node rank {NODE_RANK}") + return None + + +def test_completion(client: openai.OpenAI, + disagg_server: RemoteDisaggOpenAIServer, model_name: str): + if len(NODE_LIST) != 2: + pytest.skip("This test is only expected to run with 2 nodes") + return + if is_pytest_node(): + print(f"running test_completion on rank {RANK} node rank {NODE_RANK}") + prompt = "What is the result of 1+1? Answer in one word: " + completion = client.completions.create( + model=model_name, + prompt=prompt, + max_tokens=10, + temperature=0.0, + ) + print(f"Output: {completion.choices[0].text}") + assert completion.id is not None + message = completion.choices[0].text + assert message.startswith('2.') + + # sleep 10 seconds to ensure a successful wait_for_endpoint_ready on rank1 + time.sleep(10) + disagg_server.terminate() + + elif is_gen_node(): + # keep gen workers alive until the test ends, again we hope the NODE_LIST is ordered by NODE_RANK + url = f"http://{get_the_other_host(0)}:{DISAGG_SERVER_PORT}/health/" + wait_for_endpoint_ready(url) + wait_for_endpoint_down(url) + assert True + else: + assert True diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py new file mode 100644 index 000000000000..3d19dfd7727b --- /dev/null +++ b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py @@ -0,0 +1,223 @@ +import os +import shutil +import subprocess +import tempfile +import uuid + +import openai +import pytest + +from tensorrt_llm._utils import get_free_port +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +from ..test_llm import get_model_path +from .openai_server import RemoteDisaggOpenAIServer, RemoteOpenAIServer +from .utils import expand_slurm_nodelist, wait_for_endpoint_down, wait_for_endpoint_ready + +RANK = int(os.environ.get("SLURM_PROCID", 0)) +NODE_RANK = int(os.environ.get("SLURM_NODEID", 0)) +NODE_LIST = expand_slurm_nodelist(os.environ.get("SLURM_NODELIST", "")) +SLURM_NTASKS_PER_NODE = int(os.environ.get("SLURM_NTASKS_PER_NODE", 1)) + +# This a multi-node QA test, use a fixed port instead of finding a free port +# so that all nodes can have the same disagg server config +DISAGG_SERVER_PORT = 8000 + + +# This test is supposed to run with 2 nodes or more +def is_ctx_node(): + assert len(NODE_LIST) == 2 + return NODE_RANK == 0 + + +def is_gen_node(): + assert len(NODE_LIST) == 2 + return NODE_RANK == 1 + + +def is_disagg_node(): + return NODE_RANK == 0 + + +# The test is run on multinodes but only the first node's output is used for assertion +def is_pytest_node(): + return NODE_RANK == 0 + + +def env(): + # Remove MPI related environment variables to isolate the ctx/gen processes + # so that they will not be in the same MPI communicator, otherwise the rank and world_size may mismatch + return { + k: v + for k, v in os.environ.items() + if not ("PMI_" in k or "OMPI_" in k or "PMIX_" in k or "SLURM_" in k) + and k not in ["UCX_TLS", "UCX_NET_DEVICES"] + } + + +@pytest.fixture +def model_name(): + return "Qwen3.5-4B" + + +@pytest.fixture +def disagg_host(): + return NODE_LIST[0] + + +@pytest.fixture(params=["etcd", "http"]) +def service_discovery(request, disagg_host: str): + if request.param == "etcd": + work_dir = tempfile.mkdtemp() + data_dir = f"{work_dir}/disagg_test-etcd-{uuid.uuid4()}" + etcd = subprocess.Popen( + [ + "etcd", + "--data-dir", + data_dir, + "--listen-client-urls", + "http://0.0.0.0:2379", + "--advertise-client-urls", + f"http://{disagg_host}:2379", + ] + ) + yield etcd, f"etcd://{disagg_host}:2379" + try: + etcd.kill() + etcd.wait(timeout=10) + shutil.rmtree(data_dir) + except Exception: + pass + else: + yield None, f"http://{disagg_host}:{DISAGG_SERVER_PORT}" + + +@pytest.fixture +def disagg_cluster_config(service_discovery: tuple): + _, uri = service_discovery + return { + "cluster_uri": uri, + "cluster_name": "", + } + + +@pytest.fixture +def worker(model_name: str, disagg_cluster_config: dict): + extra_config = { + "disagg_cluster": disagg_cluster_config, + "cache_transceiver_config": {"backend": "DEFAULT"}, + "kv_cache_config": { + "free_gpu_memory_fraction": 0.5, + "enable_block_reuse": False, + }, + "disable_overlap_scheduler": True, + } + # start workers on 0.0.0.0:, then the workers should be able to + # report their correct hostname:port to the disagg server + port = get_free_port() + if is_ctx_node(): + print(f"starting ctx_server for rank {RANK} node rank {NODE_RANK}") + model_path = get_model_path(model_name) + tp_size, pp_size = 1, 1 + args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] + with RemoteOpenAIServer( + model_path, + port=port, + cli_args=args, + host="0.0.0.0", + env=env(), + llmapi_launch=False, + rank=RANK % SLURM_NTASKS_PER_NODE, + extra_config=extra_config, + role=ServerRole.CONTEXT, + ) as server: + yield server + elif is_gen_node(): + print(f"starting gen_server for rank {RANK} node rank {NODE_RANK}") + model_path = get_model_path(model_name) + tp_size, pp_size = 1, 1 + args = ["--tp_size", str(tp_size), "--pp_size", str(pp_size)] + with RemoteOpenAIServer( + model_path, + port=port, + cli_args=args, + host="0.0.0.0", + env=env(), + llmapi_launch=False, + rank=RANK % SLURM_NTASKS_PER_NODE, + extra_config=extra_config, + role=ServerRole.GENERATION, + ) as server: + yield server + else: + yield None + + +# different from non-service-discovery version, disagg server doesn't have to +# wait for ctx/gen servers to get ready +@pytest.fixture +def disagg_server(disagg_cluster_config: dict): + if is_disagg_node(): + disagg_config = { + "disagg_cluster": disagg_cluster_config, + "port": DISAGG_SERVER_PORT, + "hostname": "0.0.0.0", + } + print(f"starting disagg_server for rank {RANK} node rank {NODE_RANK}") + # ctx/gen servers are unnecessary for service discovery test + with RemoteDisaggOpenAIServer( + ctx_servers=[], + gen_servers=[], + port=DISAGG_SERVER_PORT, + disagg_config=disagg_config, + llmapi_launch=False, + env=env(), + wait_ready=False, # wait it to be ready in test body + ) as server: + yield server + else: + print(f"skipping disagg_server for rank {RANK} node rank {NODE_RANK}") + yield None + + +@pytest.fixture +def client(disagg_server: RemoteDisaggOpenAIServer): + if is_pytest_node(): + return disagg_server.get_client() + else: + print(f"skipping client for rank {RANK} node rank {NODE_RANK}") + return None + + +def test_completion( + disagg_server: RemoteDisaggOpenAIServer, + worker: RemoteOpenAIServer, + client: openai.OpenAI, + disagg_host: str, + model_name: str, +): + disagg_health_url = f"http://{disagg_host}:{DISAGG_SERVER_PORT}/health/" + wait_for_endpoint_ready(disagg_health_url) + if is_pytest_node(): + print(f"running test_completion on rank {RANK} node rank {NODE_RANK}") + prompt = "What is the result of 1+1? Answer in one word: " + for _ in range(10): + completion = client.completions.create( + model=model_name, + prompt=prompt, + max_tokens=10, + temperature=0.0, + ) + print(f"Output: {completion.choices[0].text}") + assert completion.id is not None + message = completion.choices[0].text + assert message.startswith("2.") + + disagg_server.terminate() + + elif is_gen_node(): + # keep gen workers alive until the test ends + wait_for_endpoint_down(disagg_health_url) + assert True + else: + assert True From 80b008ce84325eae6937c671cfe6455ac21598ac Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 15:45:15 +0800 Subject: [PATCH 11/16] [None][chore] unwaive fixed tests in waives.txt Remove waivers for TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm(_eagle)] on GB300 (nvbugs/6572838), test_overlap_scheduler_consistency on GB300 (nvbugs/6608387), and the MiniMax M3 MSA sparse-attention decode tests (nvbugs/6669902). Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index e6dc7b0ed37e..55cfd39cd8f3 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -158,16 +158,12 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_ full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_fp8_block_scales_4gpus[pp4-mtp_nextn=0-fp8kv=True-attention_dp=False-cuda_graph=True-overlap_scheduler=True-torch_compile=True-sampler_async_worker=False] SKIP (https://nvbugs/6661948) full:GB300/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=CUTLASS-mtp_nextn=0-pp4-fp8kv=False-attention_dp=False-cuda_graph=False-overlap_scheduler=False-low_precision_combine=False-torch_compile=True] SKIP (https://nvbugs/6697099) full:GB300/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] SKIP (https://nvbugs/6714109) -full:GB300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm] SKIP (https://nvbugs/6572838) -full:GB300/accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_nvfp4_4gpus[latency_moe_trtllm_eagle] SKIP (https://nvbugs/6572838) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_fp8_moe_dflash SKIP (https://nvbugs/6316985) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (https://nvbugs/6487918) full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_8_2_4T_A95B::test_fp8_tp16_mtp3_trtllm SKIP (https://nvbugs/6694922) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) -full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[block_reuse-python_scheduler] SKIP (https://nvbugs/6608387) -full:GB300/unittest/_torch/executor/test_overlap_scheduler.py::test_overlap_scheduler_consistency[no_reuse-python_scheduler] SKIP (https://nvbugs/6608387) full:GB300/unittest/_torch/modeling/test_modeling_gpt_oss.py::test_gpt_oss_trtllmgen[CUTLASS] SKIP (https://nvbugs/6633932) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=False] SKIP (https://nvbugs/6313072) full:H100/accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_auto_dtype[mtp_nextn=2-overlap_scheduler=True] SKIP (https://nvbugs/6313072) @@ -261,9 +257,6 @@ test_e2e.py::test_ptp_quickstart_advanced[Nemotron-Nano-9B-v2-nvfp4-NVIDIA-Nemot test_e2e.py::test_ptp_quickstart_bert[TRTLLM-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (https://nvbugs/6605819) test_e2e.py::test_ptp_quickstart_bert[VANILLA-BertForSequenceClassification-bert/bert-base-uncased-yelp-polarity] SKIP (bug pending, tracked in PR 17414) unittest/_torch/attention/sparse/dsa/test_req_idx_per_token.py::test_on_update_kv_lens_rebuilds_stale_map SKIP (https://nvbugs/6574939) -unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype0] SKIP (https://nvbugs/6669902) -unittest/_torch/attention/sparse/msa/test_minimax_m3_index_decode_score.py::test_index_decode_score_matches_msa_proxy[dtype1] SKIP (https://nvbugs/6669902) -unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py::test_sparse_decode_matches_msa_kernel SKIP (https://nvbugs/6669902) unittest/_torch/attention/sparse/msa/test_msa_backend.py::test_msa_paged_hnd_input_materializes_unaligned_outer_stride SKIP (https://nvbugs/6661846) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[exaone_moe_gqa_swa128-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6668773) unittest/_torch/attention/test_attention_backends.py::test_attention_backend[qwen2_0_5b_gqa_hd64-ctx-bf16-HND-p32-v1] SKIP (https://nvbugs/6641268) From 4879c7722cf8f48db7d67924790313a8aa2285d8 Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Sun, 13 Sep 2026 16:10:00 +0800 Subject: [PATCH 12/16] [TRTLLM-15124][test] use TestQwen3_30B_A3B in test_torch_compile_nodeids_are_private Replace the TestGemma3_1BInstruct placeholder (introduced when retiring Llama-3.1-8B coverage) with a real, currently-existing torch_compile parametrized test ID from TestQwen3_30B_A3B::test_fp8, instead of substituting an unrelated model class. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/unittest/llmapi/test_session_reuse.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/unittest/llmapi/test_session_reuse.py b/tests/unittest/llmapi/test_session_reuse.py index ee2c79c70fde..b346d1c94041 100644 --- a/tests/unittest/llmapi/test_session_reuse.py +++ b/tests/unittest/llmapi/test_session_reuse.py @@ -497,16 +497,14 @@ def test_torch_compile_nodeids_are_private(): from test_common.session_reuse_hooks import _is_private_nodeid assert _is_private_nodeid( - "accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::" - "test_fp8_prequantized[torch_compile=True]" + "accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=True]" ) assert _is_private_nodeid( "accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::" "test_nvfp4_multi_gpus_piecewise_cuda_graph[baseline]" ) assert not _is_private_nodeid( - "accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::" - "test_fp8_prequantized[torch_compile=False]" + "accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_fp8[latency-torch_compile=False]" ) assert not _is_private_nodeid( "unittest/llmapi/test_llm_args.py::test_torch_compile_config_round_trip" From 8f9eb28d7c7410686ae64f214c27b7c9273783fd Mon Sep 17 00:00:00 2001 From: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> Date: Mon, 14 Sep 2026 14:48:32 +0800 Subject: [PATCH 13/16] Update waives.txt Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 85732341990c..592e96f78c42 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -134,7 +134,6 @@ full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_a full:B300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp4_genpp4[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6771023) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp4_gentp4[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) -full:B300/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_mamba_conc_greater_than_mbs[NVIDIA-Nemotron-3-Super-120B-A12B-FP8] SKIP (https://nvbugs/6770978) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp4-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:B300/disaggregated/test_disaggregated.py::test_disaggregated_qwen3_32b_fp8[Qwen3/Qwen3-32B-FP8] SKIP (https://nvbugs/6770977) From de30ca9fa73d3bfa5777bad1f3f6415bf1d52193 Mon Sep 17 00:00:00 2001 From: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:52:04 +0800 Subject: [PATCH 14/16] Update waives.txt Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index a1dae11bb00d..ec076aeca20e 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -185,7 +185,6 @@ full:H20/accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[t full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:H20/accuracy/test_llm_api_pytorch_multimodal.py::TestNanoV3Omni::test_auto_dtype[bf16] SKIP (https://nvbugs/6618649) -full:H20/disaggregated/test_disaggregated.py::test_disaggregated_overlap_gen_first[ctx_pp1-TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6692009) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxpp2_gentp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) full:L40S/disaggregated/test_disaggregated.py::test_disaggregated_ctxtp2_genpp2[TinyLlama-1.1B-Chat-v1.0] SKIP (https://nvbugs/6728119) From da2508855e59becd8ebcf47c0f270e9564d7c8cb Mon Sep 17 00:00:00 2001 From: "Xin He (SW-GPU)" <200704525+xinhe-nv@users.noreply.github.com> Date: Wed, 16 Sep 2026 09:58:04 +0800 Subject: [PATCH 15/16] [TRTLLM-15124][fix] fix Mamba KV cache OOM in multi-node disagg completion tests Qwen3.5-4B is a hybrid Mamba/attention model whose recurrent-state + attention-page cache needs more headroom than Llama-3.1-8B did: free_gpu_memory_fraction=0.5 undersizes the V2 Mamba cache quota (observed: got ~34.5 GiB, needed ~50.3 GiB). Raise it to 0.8 in both _test_disagg_serving_multi_nodes*.py worker configs, and raise the skip_less_device_memory gate in test_e2e.py from 40000 to 80000 so the tests only run on GPUs that can actually fit it. Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/defs/test_e2e.py | 4 ++-- .../unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py | 2 +- .../_test_disagg_serving_multi_nodes_service_discovery.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 7debf2780452..39affac4c0e3 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -654,7 +654,7 @@ def test_trtllm_multimodal_benchmark_serving(llm_root, llm_venv): @skip_pre_hopper @pytest.mark.skip_less_device(4) -@pytest.mark.skip_less_device_memory(40000) +@pytest.mark.skip_less_device_memory(80000) @pytest.mark.parametrize("service_discovery", ["etcd"]) def test_openai_disagg_multi_nodes_completion_service_discovery( llm_root, llm_venv, service_discovery): @@ -670,7 +670,7 @@ def test_openai_disagg_multi_nodes_completion_service_discovery( @skip_pre_hopper @pytest.mark.skip_less_device(4) -@pytest.mark.skip_less_device_memory(40000) +@pytest.mark.skip_less_device_memory(80000) @pytest.mark.parametrize("gen_config", ["gen_tp2pp1", "gen_tp1pp2", "gen_tp1pp1"]) @pytest.mark.parametrize("ctx_config", diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py index 2324e5a5aece..74b1119f787b 100644 --- a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py +++ b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py @@ -115,7 +115,7 @@ def worker(model_name: str, ctx_tp_pp_size: tuple, gen_tp_pp_size: tuple): "backend": "DEFAULT" }, "kv_cache_config": { - "free_gpu_memory_fraction": 0.5, + "free_gpu_memory_fraction": 0.8, "enable_block_reuse": False, }, "disable_overlap_scheduler": True, diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py index 3d19dfd7727b..2d78d05dff45 100644 --- a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py +++ b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py @@ -107,7 +107,7 @@ def worker(model_name: str, disagg_cluster_config: dict): "disagg_cluster": disagg_cluster_config, "cache_transceiver_config": {"backend": "DEFAULT"}, "kv_cache_config": { - "free_gpu_memory_fraction": 0.5, + "free_gpu_memory_fraction": 0.8, "enable_block_reuse": False, }, "disable_overlap_scheduler": True, From 4f99b4b03b4fef661f24670e242b059bee4a381d Mon Sep 17 00:00:00 2001 From: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> Date: Thu, 17 Sep 2026 12:07:29 +0800 Subject: [PATCH 16/16] Update waives.txt Signed-off-by: xinhe-nv <200704525+xinhe-nv@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 1b31b196bcab..c970fa93d2a9 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -162,7 +162,6 @@ full:GB300/accuracy/test_llm_api_pytorch.py::TestQwen3_5_4B::test_bf16 SKIP (htt full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[forced_chunked_prefill] SKIP (https://nvbugs/6597570) full:GB300/accuracy/test_llm_api_pytorch_multimodal.py::TestExaone4_5_33B::test_auto_dtype[full_budget] SKIP (https://nvbugs/6597570) full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_attention_dp_overlap[DeepSeek-V3-Lite-fp8] SKIP (https://nvbugs/6581064) -full:GB300/disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct] SKIP (https://nvbugs/6275959) full:GB300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[cudagraph] SKIP (https://nvbugs/6758594) full:GB300/llmapi/test_llm_api_pytorch_moe_lora.py::test_qwen_moe_routed_expert_multi_lora_varying_ranks[eager] SKIP (https://nvbugs/6758594) full:GB300/unittest/_torch/modeling/test_modeling_gpt_oss.py::test_gpt_oss_trtllmgen[CUTLASS] SKIP (https://nvbugs/6633932)