eagle passing

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

tests/v1/spec_decode/test_eagle.py

@@ -11,6 +11,7 @@
                          VllmConfig)
 from vllm.model_executor.models.llama import LlamaForCausalLM
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.spec_decode.eagle import EagleProposer
 
 model_dir = "meta-llama/Llama-3.1-8B-Instruct"
@@ -52,6 +53,31 @@ def _create_proposer(method: str, k: int) -> EagleProposer:
                          device=current_platform.device_type)
 
 
+def _create_common_attn_metadata(
+        cu_target_query_lens: torch.Tensor,
+        device: torch.device) -> CommonAttentionMetadata:
+    """Create minimal CommonAttentionMetadata for testing."""
+    batch_size = cu_target_query_lens.shape[0] - 1
+    num_tokens = cu_target_query_lens[-1].item()
+    seq_lens = cu_target_query_lens[1:] - cu_target_query_lens[:-1]
+
+    return CommonAttentionMetadata(
+        query_start_loc=cu_target_query_lens,
+        query_start_loc_cpu=cu_target_query_lens.cpu(),
+        seq_lens=seq_lens,
+        seq_lens_cpu=seq_lens.cpu(),
+        num_computed_tokens_cpu=seq_lens.cpu(),
+        num_reqs=batch_size,
+        num_actual_tokens=int(num_tokens),
+        max_query_len=int(seq_lens.max().item()),
+        block_table_tensor=torch.zeros((batch_size, 1),
+                                       dtype=torch.int32,
+                                       device=device),
+        slot_mapping=torch.arange(num_tokens, dtype=torch.int64,
+                                  device=device),
+    )
+
+
 def test_prepare_inputs():
     """
     cu_target_query_lens: [0, a, a + b, a + b + c]
@@ -106,13 +132,19 @@ def test_prepare_inputs():
         device=device)
 
     # n1 + n2 + n3 - a - b -c
-    num_tokens = cu_target_query_lens[-1].item() - num_rejected_tokens.sum(
-    ).item()
+    num_tokens = int(cu_target_query_lens[-1].item() -
+                     num_rejected_tokens.sum().item())
 
-    cu_num_tokens, token_indices = EagleProposer.prepare_inputs(
-        cu_target_query_lens, num_rejected_tokens, num_tokens)
+    # Create CommonAttentionMetadata for new API
+    common_attn_metadata = _create_common_attn_metadata(
+        cu_target_query_lens, device)
+    proposer = _create_proposer("eagle", 1)
 
-    assert torch.equal(cu_num_tokens, expected_cu_num_tokens)
+    updated_metadata, token_indices = proposer.prepare_inputs(
+        common_attn_metadata, num_rejected_tokens.cpu(), num_tokens)
+
+    assert torch.equal(updated_metadata.query_start_loc,
+                       expected_cu_num_tokens)
     assert token_indices.shape[0] == expected_cu_num_tokens[-1].item()
     assert torch.equal(token_indices, expected_token_indices)
 
@@ -284,26 +316,33 @@ def create_deterministic_logits(token_ids):
     target_hidden_states = torch.randn(total_tokens,
                                        hidden_size,
                                        device=device)
-    target_slot_mapping = torch.randint(0,
-                                        100, (total_tokens, ),
-                                        device=device)
     next_token_ids = torch.randint(0,
                                    vocab_size, (batch_size, ),
                                    dtype=torch.int32,
                                    device=device)
-    block_table = torch.randint(0, 10, (batch_size, 10), device=device)
-
     sampling_metadata = mock.MagicMock()
 
-    # Call the method under test
-    result = proposer.propose(target_token_ids=target_token_ids,
-                              target_positions=target_positions,
-                              target_hidden_states=target_hidden_states,
-                              target_slot_mapping=target_slot_mapping,
-                              next_token_ids=next_token_ids,
-                              cu_num_tokens=cu_num_tokens,
-                              block_table=block_table,
-                              sampling_metadata=sampling_metadata)
+    # Create CommonAttentionMetadata for new API
+    common_attn_metadata = _create_common_attn_metadata(cu_num_tokens, device)
+
+    # Mock runner for attention metadata building
+    proposer.runner = mock.MagicMock()
+    proposer.runner.attn_metadata_builders = [mock.MagicMock()]
+
+    # Create mock with required attributes for multi-token tests
+    attn_metadata_mock = mock.MagicMock()
+    attn_metadata_mock.max_seq_len = 10
+    attn_metadata_mock.seq_lens = torch.tensor([5, 3], device=device)
+    proposer.runner.attn_metadata_builders[
+        0].build.return_value = attn_metadata_mock
+
+    with mock.patch('vllm.v1.spec_decode.eagle.isinstance', return_value=True):
+        result = proposer.propose(target_token_ids=target_token_ids,
+                                  target_positions=target_positions,
+                                  target_hidden_states=target_hidden_states,
+                                  next_token_ids=next_token_ids,
+                                  common_attn_metadata=common_attn_metadata,
+                                  sampling_metadata=sampling_metadata)
 
     assert result.shape == (batch_size, num_speculative_tokens)
 

vllm/v1/attention/backends/flash_attn.py

@@ -290,8 +290,8 @@ def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
             prefix_kv_lens = torch.tensor([common_prefix_len],
                                           dtype=torch.int32,
                                           device=self.device)
-            suffix_kv_lens = (seq_lens_cpu[:num_reqs] - common_prefix_len)
-            suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(self.device)
+            suffix_kv_lens = (seq_lens_cpu[:num_reqs] - common_prefix_len).to(
+                self.device, non_blocking=True)
             prefix_scheduler_metadata = schedule(
                 batch_size=1,
                 cu_query_lens=cu_prefix_query_lens,

vllm/v1/spec_decode/eagle.py

@@ -37,16 +37,14 @@ def __init__(
         self.method = self.speculative_config.method
 
         self.runner = runner
-        self.arange_np = np.arange(vllm_config.scheduler_config.max_num_seqs +
-                                   1)
-
         self.dtype = vllm_config.model_config.dtype
         self.max_model_len = vllm_config.model_config.max_model_len
         self.block_size = vllm_config.cache_config.block_size
         self.num_speculative_tokens = (
             self.speculative_config.num_speculative_tokens)
         self.max_num_tokens = (
             vllm_config.scheduler_config.max_num_batched_tokens)
+        self.arange_np = np.arange(self.max_num_tokens)
         # We need to get the hidden size from the draft model config because
         # the draft model's hidden size can be different from the target model's
         # hidden size (e.g., Llama 3.3 70B).
@@ -286,7 +284,14 @@ def prepare_inputs(
         # Step 3. [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
         arange = self.arange_np[:total_num_tokens] - cumsums_offsets
 
-        tokens_indices = arange + query_start_loc_cpu[:-1]
+        # Expand starting positions to match token pattern
+        query_start_expanded = np.repeat(query_start_loc_cpu[:-1].numpy(),
+                                         num_tokens_per_req.numpy())
+        tokens_indices = arange + query_start_expanded
+
+        # Ensure tokens_indices are within valid range for slot_mapping
+        max_slot_idx = common_attn_metadata.slot_mapping.size(0) - 1
+        tokens_indices = np.clip(tokens_indices, 0, max_slot_idx)
 
         spec_common_attn_metadata = CommonAttentionMetadata(
             query_start_loc=spec_query_start_loc_cpu.to(device,
@@ -297,13 +302,14 @@ def prepare_inputs(
             num_computed_tokens_cpu=(
                 common_attn_metadata.num_computed_tokens_cpu),
             num_reqs=common_attn_metadata.num_reqs,
-            num_actual_tokens=num_tokens,
+            num_actual_tokens=total_num_tokens,
             max_query_len=query_len_per_req.max().item(),
             block_table_tensor=common_attn_metadata.block_table_tensor,
             slot_mapping=common_attn_metadata.slot_mapping[tokens_indices],
         )
 
-        return spec_common_attn_metadata, tokens_indices
+        return spec_common_attn_metadata, torch.from_numpy(tokens_indices).to(
+            device)
 
     def load_model(self, target_model: nn.Module) -> None:
         draft_model_config = \