Swapping the order of sampling operations in the conditional selector. (#199)

Adding P3L measurement to the benchmarks collection tools. A more beautiful version of the code with "Swapping the order of sampling operations in the conditional selector. (#199)"

Author: Alexei-V-Ivanov-AMD

vllm/model_executor/layers/sampler.py

@@ -750,16 +750,16 @@ def get_pythonized_sample_results(
         if sampling_type not in sample_metadata:
             continue
         (seq_group_id, seq_groups) = sample_metadata[sampling_type]
-        if sampling_type == SamplingType.FORCED:
-            sample_results = _forced_sample(seq_groups, forced_samples)
-        elif sampling_type == SamplingType.GREEDY:
+        if sampling_type == SamplingType.GREEDY:
             sample_results = _greedy_sample(seq_groups, greedy_samples)
         elif sampling_type in (SamplingType.RANDOM, SamplingType.RANDOM_SEED):
             sample_results = _random_sample(seq_groups,
                                             multinomial_samples[sampling_type])
         elif sampling_type == SamplingType.BEAM:
             sample_results = _beam_search_sample(seq_groups,
                                                  beam_search_logprobs)
+        elif sampling_type == SamplingType.FORCED:
+            sample_results = _forced_sample(seq_groups, forced_samples)
         sample_results_dict.update(zip(seq_group_id, sample_results))
 
     return [
@@ -825,19 +825,8 @@ def _sample_with_torch(
         seq_groups = [sampling_metadata.seq_groups[i] for i in seq_group_id]
         sample_metadata[sampling_type] = (seq_group_id, seq_groups)
         long_sample_indices = sample_indices.long()
-        if sampling_type == SamplingType.FORCED:
-            if (seq_groups[0].sampling_params.future_context is not None):
-                forced_samples = torch.tensor([
-                    seq_groups[0].sampling_params.future_context[0][min(
-                        len(sampling_metadata.seq_groups[0].seq_data[
-                            sampling_params.cntr].output_token_ids),
-                        len(seq_groups[0].sampling_params.future_context[0]) -
-                        1)]
-                ])
-            else:
-                forced_samples = torch.argmax(logprobs[long_sample_indices],
-                                              dim=-1)
-        elif sampling_type == SamplingType.GREEDY:
+
+        if sampling_type == SamplingType.GREEDY:
             greedy_samples = torch.argmax(logprobs[long_sample_indices],
                                           dim=-1)
 
@@ -886,6 +875,18 @@ def _sample_with_torch(
 
         elif sampling_type == SamplingType.BEAM:
             beam_search_logprobs = logprobs[sample_indices]
+        elif sampling_type == SamplingType.FORCED:
+            if (seq_groups[0].sampling_params.future_context is not None):
+                forced_samples = torch.tensor([
+                    seq_groups[0].sampling_params.future_context[0][min(
+                        len(sampling_metadata.seq_groups[0].seq_data[
+                            sampling_params.cntr].output_token_ids),
+                        len(seq_groups[0].sampling_params.future_context[0]) -
+                        1)]
+                ])
+            else:
+                forced_samples = torch.argmax(logprobs[long_sample_indices],
+                                              dim=-1)
         else:
             raise ValueError(f"Unsupported sampling type: {sampling_type}")