From 2c5639deac66a0727e9949a23920fc3d6ab27163 Mon Sep 17 00:00:00 2001 From: Izzy Putterman Date: Tue, 23 Sep 2025 15:50:05 -0700 Subject: [PATCH] Eagle: Attention DP Signed-off-by: Izzy Putterman --- .../_torch/models/modeling_speculative.py | 29 +++++++---- .../_torch/speculative/test_eagle3.py | 51 ++++++++++--------- 2 files changed, 47 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 6eb989af33fa..31d52791f6b1 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -48,6 +48,8 @@ def __init__( ) tp_size = model_config.mapping.tp_size + if model_config.mapping.enable_attention_dp: + tp_size = 1 # Override the QKV projection. The number of input features # is twice as big for EAGLE3 draft models. self.qkv_proj = Linear( @@ -55,7 +57,7 @@ def __init__( tp_size * self.q_size + 2 * tp_size * self.kv_size, bias=config.attention_bias, dtype=config.torch_dtype, - mapping=model_config.mapping, + mapping=self.qkv_proj.mapping, tensor_parallel_mode=TensorParallelMode.COLUMN, weights_loading_config=WeightsLoadingConfig( weight_mode=WeightMode.FUSED_QKV_LINEAR), @@ -89,6 +91,8 @@ def __init__( bias=getattr(config, "mlp_bias", False), dtype=config.torch_dtype, config=model_config, + overridden_tp_size=1 + if model_config.mapping.enable_attention_dp else None, ) self.input_layernorm = RMSNorm(hidden_size=config.hidden_size, eps=config.rms_norm_eps, @@ -182,14 +186,21 @@ def __init__( requires_grad=False) if self.hidden_size_in != config.hidden_size: - self.embed_tokens = Embedding( - config.vocab_size, - config.hidden_size, - dtype=config.torch_dtype, - mapping=model_config.mapping, - tensor_parallel_mode=TensorParallelMode.COLUMN, - gather_output=True, - ) + if model_config.mapping.enable_attention_dp: + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + ) + else: + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) else: # Shared with target model. self.embed_tokens = None diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index a37ce3181994..e32f715c50f9 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -24,38 +24,40 @@ def enforce_single_worker(monkeypatch): @pytest.mark.parametrize( - "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch", + "use_cuda_graph,attn_backend,disable_overlap_scheduler,enable_block_reuse,use_one_model,enable_chunked_prefill,use_chain_drafter,multi_batch,attention_dp", [ - [True, "TRTLLM", True, False, False, False, True, False], - [True, "TRTLLM", True, False, False, False, False, False], - [False, "TRTLLM", True, False, False, False, True, False], - [False, "TRTLLM", True, False, False, False, False, False], - [True, "FLASHINFER", True, False, False, False, True, False], - [False, "FLASHINFER", True, False, False, False, True, False], - [False, "TRTLLM", False, True, True, False, True, False], - [True, "TRTLLM", False, True, True, False, True, False], - [True, "TRTLLM", True, False, True, True, True, False], - [True, "TRTLLM", True, False, True, False, True, False], + [True, "TRTLLM", True, False, False, False, True, False, False], + [True, "TRTLLM", True, False, False, False, False, False, False], + [False, "TRTLLM", True, False, False, False, True, False, False], + [False, "TRTLLM", True, False, False, False, False, False, False], + [True, "FLASHINFER", True, False, False, False, True, False, False], + [False, "FLASHINFER", True, False, False, False, True, False, False], + [False, "TRTLLM", False, True, True, False, True, False, False], + [True, "TRTLLM", False, True, True, False, True, False, False], + [True, "TRTLLM", True, False, True, True, True, False, False], + [True, "TRTLLM", True, False, True, False, True, False, False], # TODO: nvbugs/5461761 # [True, "TRTLLM", True, False, False, True, True, False], - [True, "TRTLLM", False, False, False, False, True, False], - [False, "TRTLLM", False, False, False, False, True, False], - [True, "TRTLLM", False, False, False, False, False, True], - [False, "TRTLLM", False, False, False, False, False, True], - [True, "TRTLLM", False, False, False, False, True, True], - [False, "TRTLLM", False, False, False, False, True, True], - [True, "TRTLLM", False, False, False, False, False, False], - [False, "TRTLLM", False, False, False, False, False, False], - [True, "TRTLLM", False, False, False, True, True, False], - [True, "TRTLLM", False, False, False, True, False, False], - [True, "FLASHINFER", False, False, False, False, True, False], - [False, "FLASHINFER", False, False, False, False, True, False], + [True, "TRTLLM", False, False, False, False, True, False, False], + [False, "TRTLLM", False, False, False, False, True, False, False], + [True, "TRTLLM", False, False, False, False, False, True, False], + [True, "TRTLLM", False, False, False, False, False, True, True], + [False, "TRTLLM", False, False, False, False, False, True, False], + [True, "TRTLLM", False, False, False, False, True, True, False], + [False, "TRTLLM", False, False, False, False, True, True, False], + [True, "TRTLLM", False, False, False, False, False, False, False], + [False, "TRTLLM", False, False, False, False, False, False, False], + [True, "TRTLLM", False, False, False, True, True, False, False], + [True, "TRTLLM", False, False, False, True, False, False, False], + [True, "FLASHINFER", False, False, False, False, True, False, False], + [False, "FLASHINFER", False, False, False, False, True, False, False], ]) @pytest.mark.high_cuda_memory 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, - use_chain_drafter: bool, multi_batch: bool, request): + use_chain_drafter: bool, multi_batch: bool, + attention_dp: bool, request): # Use enforce_single_worker fixture only when use_chain_drafter is False. # Otherwise, we can't modify the returned value of _get_allow_chain_drafter in multiprocessing. if not use_chain_drafter: @@ -98,6 +100,7 @@ def test_llama_eagle3(use_cuda_graph: bool, attn_backend: str, 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