@@ -353,71 +353,6 @@ def _safely_apply_temperature_inplace(
353353 return logits_inout .div_ (safe_temp .unsqueeze (dim = 1 ))
354354
355355
356- def _apply_top_k_top_p (
357- logits : torch .Tensor ,
358- k : Optional [torch .Tensor ],
359- p : Optional [torch .Tensor ],
360- ) -> torch .Tensor :
361- logits_sort , logits_idx = logits .sort (dim = - 1 , descending = False )
362- if k is not None :
363- top_k_mask = logits_sort .size (1 ) - k .to (torch .long )
364- top_k_mask = top_k_mask .clamp (min = 0 )
365- top_k_mask = logits_sort .gather (1 , top_k_mask .unsqueeze (dim = 1 ))
366- top_k_mask = logits_sort < top_k_mask
367- logits_sort .masked_fill_ (top_k_mask , - float ("inf" ))
368- if p is not None :
369- probs_sort = logits_sort .softmax (dim = - 1 )
370- probs_sum = torch .cumsum (probs_sort , dim = - 1 , out = probs_sort )
371- top_p_mask = probs_sum <= 1 - p .unsqueeze (dim = 1 )
372- top_p_mask [:, - 1 ] = False
373- logits_sort .masked_fill_ (top_p_mask , - float ("inf" ))
374- return logits_sort .scatter (dim = - 1 , index = logits_idx , src = logits_sort )
375-
376-
377- def _random_sample (probs : torch .Tensor ) -> torch .Tensor :
378- q = torch .empty_like (probs ).exponential_ ()
379- return probs .div_ (q ).argmax (dim = - 1 ).view (- 1 )
380-
381-
382- def forward_native_sampling (
383- logits : torch .Tensor ,
384- k : Optional [torch .Tensor ],
385- p : Optional [torch .Tensor ],
386- ) -> torch .Tensor :
387- logits = _apply_top_k_top_p (logits , k , p )
388- probs = logits .softmax (dim = - 1 , dtype = torch .float32 )
389- return _random_sample (probs )
390-
391-
392- def compute_probs_from_logits_op (
393- logits : torch .Tensor ,
394- temperatures : torch .Tensor ,
395- top_k : Optional [torch .Tensor ],
396- top_p : Optional [torch .Tensor ],
397- ) -> torch .Tensor :
398- """Pure-PyTorch CPU fallback for probability computation."""
399- is_greedy = temperatures <= _GREEDY_TEMPERATURE_THRESHOLD
400- # Greedy rows must pick the argmax of the *original* logits (before temperature).
401- # Capture the argmax up front; _safely_apply_temperature_inplace then guards the
402- # division against the greedy sentinel.
403- argmax_ids = logits .argmax (dim = - 1 , keepdim = True )
404-
405- logits = _safely_apply_temperature_inplace (logits , temperatures )
406- logits = _apply_top_k_top_p (logits , top_k , top_p )
407- probs = logits .softmax (dim = - 1 , dtype = torch .float32 )
408-
409- # Turn the greedy rows into a one-hot at argmax by editing `probs` in place,
410- # instead of building a full [batch, vocab] one-hot buffer and a [batch, vocab]
411- # torch.where copy. The torch.where here only runs on a [batch, 1] tensor.
412- # NB: argwhere/index-select on the greedy rows would give a data-dependent shape
413- # that breaks the surrounding torch.compile graph, so we keep it dense.
414- greedy_col = is_greedy .unsqueeze (1 )
415- new_at_argmax = torch .where (greedy_col , 1.0 , probs .gather (1 , argmax_ids ))
416- probs .masked_fill_ (greedy_col , 0.0 )
417- probs .scatter_ (1 , argmax_ids , new_at_argmax )
418- return probs
419-
420-
421356class _Fusions :
422357 @staticmethod
423358 @torch .compile (dynamic = None , fullgraph = True )
0 commit comments