zero temp sampling (pytorch#277)

Add special case for zero-temperature sampling. For stories15M on my devserver, seems to improve tokens/sec as follows:

before: 189, 180, 166
after: 264, 285, 285

Author: swolchok

generate.py

@@ -101,7 +101,11 @@ def logits_to_probs(logits, temperature: float = 1.0, top_k: Optional[int] = Non
     return probs
 
 
-def sample(logits, temperature: float = 1.0, top_k: Optional[int] = None):
+def sample(logits, need_probs: bool, temperature: float = 1.0,top_k: Optional[int] = None):
+    if temperature == 0 and not need_probs:
+        _, idx_next = torch.topk(logits, k=1, dim=-1)
+        idx_next = idx_next.squeeze(dim=(0, 1))
+        return (idx_next, None)
     probs = logits_to_probs(logits[0, -1], temperature, top_k)
     idx_next = multinomial_sample_one_no_sync(probs)
     return idx_next, probs
@@ -129,23 +133,24 @@ def prefill(
         # input_pos: [B, S]
         logits = model(x, input_pos)
 
-    return sample(logits, **sampling_kwargs)[0]
+    return sample(logits, need_probs=False, **sampling_kwargs)[0]
 
 
 def decode_one_token(
-    model: Transformer, x: torch.Tensor, input_pos: torch.Tensor, **sampling_kwargs
-) -> Tuple[torch.Tensor, torch.Tensor]:
+    model: Transformer, x: torch.Tensor, input_pos: torch.Tensor, need_probs: bool, **sampling_kwargs
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
     # input_pos: [B, 1]
     assert input_pos.shape[-1] == 1
     logits = model(x, input_pos)
-    return sample(logits, **sampling_kwargs)
+    return sample(logits, need_probs=need_probs, **sampling_kwargs)
 
 
 def decode_n_tokens(
     model: Transformer,
     cur_token: torch.Tensor,
     input_pos: torch.Tensor,
     num_new_tokens: int,
+    need_probs: bool,
     callback=lambda _: _,
     **sampling_kwargs,
 ):
@@ -154,12 +159,13 @@ def decode_n_tokens(
         # Actually better for Inductor to codegen attention here
         with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]):
             next_token, next_prob = decode_one_token(
-                model, cur_token, input_pos, **sampling_kwargs
+                model, cur_token, input_pos, need_probs=need_probs, **sampling_kwargs
             )
             input_pos += 1
             new_tokens.append(next_token.clone())
             callback(new_tokens[-1])
-            new_probs.append(next_prob.clone())
+            if need_probs:
+                new_probs.append(next_prob.clone())
             cur_token = next_token.view(1, -1)
 
     return new_tokens, new_probs
@@ -187,6 +193,7 @@ def speculative_decode(
         cur_token.view(1, -1),
         orig_input_pos.clone(),
         speculate_k,
+        need_probs=True,
         **sampling_kwargs,
     )
 
@@ -301,6 +308,7 @@ def generate(
             input_pos,
             max_new_tokens - 1,
             callback=callback,
+            need_probs = False,
             **sampling_kwargs,
         )
         seq[T + 1 :] = torch.cat(generated_tokens)