[NPU] Avoid device synchronization in Ascend sampling (#39404)
This commit is contained in:
@@ -492,18 +492,21 @@ class TestMergeBatch(CustomTestCase):
|
||||
need_top_p_sampling=False,
|
||||
need_top_k_sampling=False,
|
||||
need_min_p_sampling=False,
|
||||
npu_top_k_top_p_eligible=True,
|
||||
)
|
||||
info2 = _make_info(
|
||||
is_all_greedy=False,
|
||||
need_top_p_sampling=True,
|
||||
need_top_k_sampling=True,
|
||||
need_min_p_sampling=True,
|
||||
npu_top_k_top_p_eligible=False,
|
||||
)
|
||||
info1.merge_batch(info2)
|
||||
self.assertFalse(info1.is_all_greedy) # AND semantics
|
||||
self.assertTrue(info1.need_top_p_sampling) # OR semantics
|
||||
self.assertTrue(info1.need_top_k_sampling) # OR semantics
|
||||
self.assertTrue(info1.need_min_p_sampling) # OR semantics
|
||||
self.assertFalse(info1.npu_top_k_top_p_eligible) # AND semantics
|
||||
|
||||
def test_merge_with_logit_bias(self):
|
||||
"""Test that merge pads missing logit_bias with zeros before concatenation."""
|
||||
@@ -675,6 +678,22 @@ class TestFromScheduleBatch(CustomTestCase):
|
||||
self.assertTrue(info.need_min_p_sampling) # 0.1 > 0
|
||||
self.assertFalse(info.is_all_greedy) # top_k=50 > 1
|
||||
|
||||
def test_npu_top_k_top_p_eligibility_uses_request_params(self):
|
||||
cases = (
|
||||
((1, 1024), True),
|
||||
((1, 1025), False),
|
||||
((4, TOP_K_ALL), False),
|
||||
)
|
||||
for top_ks, expected in cases:
|
||||
with self.subTest(top_ks=top_ks):
|
||||
batch = MagicMock()
|
||||
batch.reqs = [self._make_req(top_k=top_k) for top_k in top_ks]
|
||||
batch.device = DEVICE
|
||||
|
||||
info = SamplingBatchInfo.from_schedule_batch(batch, VOCAB_SIZE)
|
||||
|
||||
self.assertEqual(info.npu_top_k_top_p_eligible, expected)
|
||||
|
||||
def test_no_logit_bias_when_all_none(self):
|
||||
"""Test that logit_bias stays None when no request has logit_bias set."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user