Revert "[Kernel] Fuse temperature + softmax in sampling for decode speedup" (#22046)

python/sglang/srt/layers/sampler.py

@@ -18,7 +18,6 @@ from sglang.srt.sampling.sampling_params import TOP_K_ALL
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils.common import crash_on_warnings, get_bool_env_var, is_cuda, is_npu
 
-_use_fused_sampling = False
 if is_cuda():
     from flashinfer.sampling import (
         min_p_sampling_from_probs,
@@ -28,15 +27,6 @@ if is_cuda():
         top_k_renorm_prob,
         top_p_renorm_prob,
     )
-
-    from sglang.srt.layers.fused_sampling import fused_temperature_softmax_inplace
-
-    _use_fused_sampling = True
-
-# Batch size threshold for fused Triton kernel vs PyTorch softmax.
-# Below this threshold, PyTorch's native div+softmax is faster.
-# At and above this threshold, the fused Triton kernel wins.
-_FUSED_SAMPLING_BATCH_THRESHOLD = 128
 if is_npu():
     import torch_npu
 
@@ -162,20 +152,11 @@ class Sampler(nn.Module):
                     logprobs = logprobs_via_logsoftmax_kernel
             else:
                 # Standard path: do softmax and sample from probs.
-                # Use fused Triton kernel for large batches where it excels;
-                # fall back to PyTorch for small batches where launch overhead dominates.
-                if (
-                    _use_fused_sampling
-                    and logits.shape[0] >= _FUSED_SAMPLING_BATCH_THRESHOLD
-                ):
-                    fused_temperature_softmax_inplace(
-                        logits, sampling_info.temperatures
-                    )
-                    probs = logits
-                else:
-                    logits.div_(sampling_info.temperatures)
-                    logits[:] = torch.softmax(logits, dim=-1)
-                    probs = logits
+                logits.div_(sampling_info.temperatures)
+
+                # In-place op to save memory
+                logits[:] = torch.softmax(logits, dim=-1)
+                probs = logits
 
                 batch_next_token_ids = self._sample_from_probs(
                     probs, sampling_info, positions, simple_sampling_case