fix(core): fix handling logprobs when use rbln_sampler (#184)

* fix indexing logprobs

* fix comment

Author: rebel-eunji

vllm_rbln/v1/sample/rbln_sampler.py

@@ -241,6 +241,8 @@ def get_bucket_sizes(max_num_seqs: int) -> list[int]:
         [1, 2, 4] + list(range(8, 256, 8)) + list(
             range(256, max_num_seqs + 1, 16))
         """
+        # FIXME(eunji.lee)
+        # Not used. To be removed.
         bucket_sizes = [i for i in [1, 2, 4] if i <= max_num_seqs]
         if max_num_seqs >= 8:
             # Step size 8 for small batch sizes, up to 256(not included)

vllm_rbln/v1/worker/optimum_model_runner.py

@@ -321,7 +321,10 @@ def execute_model(
                     sampler_output.sampled_token_ids[:num_reqs]
                 if sampler_output.logprobs_tensors is not None:
                     sampler_output.logprobs_tensors = \
-                        sampler_output.logprobs_tensors[:num_reqs]
+                        self.post_process_logprobs_tensors(
+                            sampler_output.logprobs_tensors,
+                            num_reqs
+                        )
 
         with record_function_or_nullcontext("Bookkeep"):
             (
@@ -1223,3 +1226,15 @@ def get_bucket_sizes(max_num_seqs: int) -> list[int]:
             # Step size 16 for larger batch sizes
             bucket_sizes += list(range(256, max_num_seqs + 1, 16))
         return bucket_sizes
+
+    def post_process_logprobs_tensors(self, logprobs_tensors: LogprobsTensors,
+                                      num_reqs: int) -> LogprobsTensors:
+        # NOTE(eunji.lee):
+        # This implementation is not efficient but kept for debugging purposes.
+        # TODO: Modify this code in the next version when the shape of
+        # logprobs_tensors changes.
+        dict = {}
+        for field_name in logprobs_tensors._fields:
+            tensor = getattr(logprobs_tensors, field_name)
+            dict[field_name] = tensor[:num_reqs]
+        return LogprobsTensors(**dict)