-
Notifications
You must be signed in to change notification settings - Fork 604
[TEST] Add eagle proposer ut #4447
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
There was a problem hiding this 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.
| 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]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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]>
What this PR does / why we need it?
Add eagle proposer ut
Does this PR introduce any user-facing change?
No