A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 02:30 UTC
[Model Runner V2] Use FP32 for Gumbel Noise (#34854)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 18, 2026 at 17:07 UTC
5fcb0cdd68ded0d1c988121d4f744c2424522259
1 file changed
+8
-9
vllm/v1/worker/gpu/sample/gumbel.py
+8
-9
@@ -85,10 +85,10 @@ def _gumbel_sample_kernel(
85
pos = tl.load(pos_ptr + batch_idx)
86
gumbel_seed = tl.randint(seed, pos)
87
88
- # Generate gumbel noise.
89
- r = tl.rand(gumbel_seed, block).to(tl.float64)
90
- gumbel_noise = -tl.log(-tl.log(r + 1e-20) + 1e-20)
91
- gumbel_noise = gumbel_noise.to(tl.float32)
88
+ # Generate gumbel noise in FP32.
89
+ u = tl.rand(gumbel_seed, block)
90
+ u = tl.maximum(u, 1e-7)
91
+ gumbel_noise = -tl.log(-tl.log(u))
92
93
# Apply temperature.
94
if APPLY_TEMPERATURE:
@@ -99,18 +99,17 @@ def _gumbel_sample_kernel(
99
# Apply gumbel noise.
100
logits = tl.where(mask, logits + gumbel_noise, float("-inf"))
101
102
- idx = tl.argmax(logits, axis=0)
102
+ value, idx = tl.max(logits, axis=0, return_indices=True)
103
token_id = block_idx * BLOCK_SIZE + idx
104
- value = tl.max(logits, axis=0)
104
tl.store(local_argmax_ptr + batch_idx * local_argmax_stride + block_idx, token_id)
105
tl.store(local_max_ptr + batch_idx * local_max_stride + block_idx, value)
106
107
108
def gumbel_sample(
109
logits: torch.Tensor, # [num_reqs, vocab_size]
111
- idx_mapping: torch.Tensor, # [num_reqs]
112
- temperature: torch.Tensor, # [num_reqs]
113
- seed: torch.Tensor, # [num_reqs]
110
+ idx_mapping: torch.Tensor, # [max_num_reqs]
111
+ temperature: torch.Tensor, # [max_num_reqs]
112
+ seed: torch.Tensor, # [max_num_reqs]
113
pos: torch.Tensor, # [num_reqs]
114
apply_temperature: bool,
115
) -> torch.Tensor: