spec decode: match the original rank computation impl for spec decodi…

…ng (#954)

aphrodite/spec_decode/util.py

@@ -43,8 +43,8 @@ def get_sampled_token_logprobs(
                                        sampled_token_ids, ]
     expanded_selected_logprobs = selected_logprobs.unsqueeze(-1).expand(
         -1, -1, vocab_size)
-    sampled_token_ids_ranks = (logprob_tensor >=
-                               expanded_selected_logprobs).sum(-1)
+    sampled_token_ids_ranks = (logprob_tensor >
+                               expanded_selected_logprobs).sum(-1).add_(1)
 
     return sampled_token_ids_ranks, selected_logprobs
 

tests/spec_decode/test_utils.py

@@ -5,9 +5,11 @@
 
 from aphrodite.common.sequence import SequenceGroupMetadata, get_all_seq_ids
 from aphrodite.modeling.layers.rejection_sampler import RejectionSampler
+from aphrodite.modeling.layers.sampler import _get_ranks
 from aphrodite.modeling.layers.typical_acceptance_sampler import (
     TypicalAcceptanceSampler)
-from aphrodite.spec_decode.util import split_batch_by_proposal_len
+from aphrodite.spec_decode.util import (get_sampled_token_logprobs,
+                                        split_batch_by_proposal_len)
 
 
 def test_get_all_seq_ids():
@@ -126,3 +128,18 @@ def mock_spec_decode_sampler(acceptance_sampler_method):
         return sampler
     else:
         raise ValueError(f"Invalid sampler name {acceptance_sampler_method}")
+
+
+def test_get_sampled_token_logprobs():
+    """Verify get_sampled_token_logprobs returns consistent rankings 
+    with regular get_ranks when probabilities match exactly.
+    """
+    logprob_tensor = torch.tensor(
+        [[[-.1, -.1]] * 2])  # shape (num_steps, batch_size, vocab_size)
+    sampled_token_tensor = torch.tensor([[1,
+                                          0]])  # shape (num_steps, batch_size)
+    ranks_spec_dec, _ = get_sampled_token_logprobs(logprob_tensor,
+                                                   sampled_token_tensor)
+    ranks_regular = _get_ranks(logprob_tensor.reshape((2, -1)),
+                               sampled_token_tensor.reshape(-1))
+    assert torch.equal(ranks_spec_dec.reshape(-1), ranks_regular)