diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 2305c0852664..c109050783f3 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -2865,6 +2865,13 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, if moe_backend == "TRTLLM" and sm_version in (120, 121): pytest.skip(f"{moe_backend} backend does not support SM 120 or 121") + # Wrong MTP draft tokens cannot be caught by the accuracy scores: + # rejected drafts are regenerated by the target model, so accuracy + # stays intact and the only observable symptom is a collapsed + # acceptance length (https://nvbugs/6460072). Assert on it for the + # ADP + LM-head-TP cases, whose draft sampling takes a dedicated path. + check_acceptance = mtp_nextn > 0 and enable_lm_head_tp_in_adp + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.70) pytorch_config = dict( disable_overlap_scheduler=not overlap_scheduler, @@ -2886,6 +2893,7 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, **pytorch_config, enable_attention_dp=attention_dp, enable_lm_head_tp_in_adp=enable_lm_head_tp_in_adp, + enable_iter_perf_stats=check_acceptance, speculative_config=mtp_config) as llm: assert llm.args.moe_config.backend == moe_backend @@ -2903,6 +2911,22 @@ def test_nvfp4_multi_gpus(self, tp_size, pp_size, ep_size, mtp_nextn, fp8kv, # task.evaluate(llm, # extra_evaluator_kwargs=dict(apply_chat_template=True)) + if check_acceptance: + stats = llm.get_stats(timeout=2) + spec_iters = [ + stat['specDecodingStats'] for stat in stats + if stat.get('specDecodingStats') + and stat['specDecodingStats']['numDraftTokens'] > 0 + ] + assert spec_iters, \ + "No iterations with speculative decoding stats" + acceptance_length = sum(s['acceptanceLength'] + for s in spec_iters) / len(spec_iters) + assert acceptance_length >= 1.5, ( + f"MTP acceptance length {acceptance_length:.2f} < 1.5: " + "draft tokens are likely wrong in the ADP + LM-head-TP " + "sampling path.") + import gc gc.collect() torch.cuda.empty_cache()