fix bugs in sampler

colossalai/inference/core/request_handler.py

@@ -180,9 +180,9 @@ class RequestHandler:
         """
         # do logit processor
         # NOTE: need to decide the granularity to process logits (sequence or batch)
-        for type in ["top_p", "top_k", "min_p"]:
+        for type in ["top_k", "top_p", "min_p"]:
             config_dict = generation_config.to_dict()
-            if type in config_dict:
+            if type in config_dict and config_dict[type] is not None:
                 logits = logit_processor(type, logits, config_dict[type])
 
         torch.cuda.synchronize()

colossalai/inference/sampler.py

@@ -21,7 +21,7 @@ def multinomial_sample(
     """
     Sample tokens in a random phase.
     """
-    random_results = torch.multinomial(probs, num_samples=1, replacement=True).cpu()
+    random_results = torch.multinomial(probs, num_samples=1).squeeze(1)
     return random_results