Skip to content

Conversation

@GDzhu01
Copy link
Contributor

@GDzhu01 GDzhu01 commented Nov 26, 2025

What this PR does / why we need it?

Add eagle proposer ut

Does this PR introduce any user-facing change?

No

@github-actions
Copy link

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

  • A PR should do only one thing, smaller PRs enable faster reviews.
  • Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
  • Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds unit tests for the EagleProposer. The tests cover initialization, model loading, dummy runs, and token generation logic. The test structure is generally good, but I found a critical issue in one of the helper method tests which is ineffective as it mocks the method it is supposed to be testing. I've provided a suggestion to fix this to ensure proper test coverage.

Comment on lines +299 to +312
def test_prepare_inputs(self):
self.proposer.token_arange_np = np.arange(10)
mock_attn = MagicMock()
mock_attn.slot_mapping = torch.tensor([0, 1, 2, 3, 4, 5])
num_rejected = torch.tensor([1, 0, 1], device=self.device)

with patch.object(self.proposer,
'_prepare_inputs',
return_value=(torch.tensor([0, 2, 5]),
torch.tensor([1, 2, 4]))):
cu_num_tokens, indices = self.proposer._prepare_inputs(
mock_attn, num_rejected)
self.assertEqual(cu_num_tokens.tolist(), [0, 2, 5])
self.assertEqual(indices.tolist(), [1, 2, 4])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The current test for _prepare_inputs is not effective as it patches the very method it intends to test. This means the test only verifies that the patch works, not the actual implementation of _prepare_inputs. This gives a false sense of test coverage for a critical piece of logic.

I've provided a corrected version of the test that invokes the real _prepare_inputs method and validates its output against expected values for a concrete scenario. This ensures the logic for calculating new token counts and indices after rejections is properly tested.

    def test_prepare_inputs(self):
        # This test verifies the logic of `_prepare_inputs` which calculates
        # the new cumulative token counts and indices for tokens that are not
        # rejected.
        # Original query lengths: [2, 5, 3] -> query_start_loc: [0, 2, 7, 10]
        # Rejected tokens per request: [1, 2, 0]
        # Expected new query lengths: [1, 3, 3]
        # Expected new cu_num_tokens: [0, 1, 4, 7]
        # Expected indices of non-rejected tokens: [0, 2, 3, 4, 7, 8, 9]
        mock_attn_metadata = MagicMock()
        mock_attn_metadata.query_start_loc = torch.tensor([0, 2, 7, 10],
                                                          device=self.device,
                                                          dtype=torch.int32)
        num_rejected = torch.tensor([1, 2, 0],
                                    device=self.device,
                                    dtype=torch.int32)

        cu_num_tokens, indices = self.proposer._prepare_inputs(
            mock_attn_metadata, num_rejected)

        expected_cu_num_tokens = torch.tensor([0, 1, 4, 7],
                                              device=self.device,
                                              dtype=torch.int32)
        expected_indices = torch.tensor([0, 2, 3, 4, 7, 8, 9],
                                        device=self.device,
                                        dtype=torch.int64)

        self.assertTrue(torch.equal(cu_num_tokens, expected_cu_num_tokens))
        self.assertTrue(torch.equal(indices, expected_indices))

Signed-off-by: GDzhu01 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant