diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index c6699b8bc5ab..862777d4f63e 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -2093,7 +2093,45 @@ def test_nvfp4_multi_gpus_chunked_prefill(self, tp_size, pp_size, ep_size, assert llm.args.moe_config.backend == moe_backend assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 - @skip_pre_blackwell + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + def test_nvfp4_multi_gpus_corner_case(self): + """ + This test is used to test the corner case of the NVFP4 model. + When using the same value for max_seq_len and max_num_tokens, there will be no + enough kv block for the dummy requests in CUDA graph warmup when creating + the py_executor before estimating kv cache. Then CUDA graph capture will be + triggered when estimating kv cache. This may cause some errors. + More info in https://nvbugs/5485325. + """ + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.80, + dtype="fp8", + enable_block_reuse=False) + pytorch_config = dict(disable_overlap_scheduler=False, + cuda_graph_config=CudaGraphConfig( + enable_padding=True, max_batch_size=1024), + moe_config=MoeConfig(backend="TRTLLM")) + + mtp_config = MTPDecodingConfig(num_nextn_predict_layers=1) + with LLM(f"{llm_models_root()}/DeepSeek-R1/DeepSeek-R1-FP4", + tensor_parallel_size=8, + pipeline_parallel_size=1, + moe_expert_parallel_size=8, + kv_cache_config=kv_cache_config, + **pytorch_config, + enable_attention_dp=False, + speculative_config=mtp_config, + max_seq_len=5120, + max_num_tokens=5120) as llm: + + assert llm.args.quant_config.quant_algo == QuantAlgo.NVFP4 + + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + def test_nvfp4_multi_gpus_corner_case(self): """ This test is used to test the corner case of the NVFP4 model. diff --git a/tests/integration/defs/test_e2e.py b/tests/integration/defs/test_e2e.py index 4ab3e69c8955..df54cd3e4568 100644 --- a/tests/integration/defs/test_e2e.py +++ b/tests/integration/defs/test_e2e.py @@ -2303,6 +2303,7 @@ def test_ptp_quickstart_advanced_8gpus(llm_root, llm_venv, model_name, f"{llm_models_root()}/{model_path}", "--tp_size=8", "--max_batch_size=32", + "--max_num_tokens=256", ], stdout=running_log) if model_name in mapping: @@ -2366,6 +2367,7 @@ def test_ptp_quickstart_advanced_2gpus_sm120(llm_root, llm_venv, model_name, "--model_dir", f"{llm_models_root()}/{model_path}", "--tp_size=2", + "--max_num_tokens=256", ])