[Perf] Optimize `clamp` to `clamp_` (#48143)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
Jul 16, 2026 at 18:41 UTC
f17be06fbe42bab61522868a78d044fd0d2e3547
7 files changed
+48
-32
vllm/model_executor/layers/attention/mla_attention.py
+4
-4
@@ -1500,7 +1500,8 @@ def build_mla_chunked_context_metadata(
1500
chunk_ends = torch.min(
1501
context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
1502
)
1503
- chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
1503
+ chunk_seq_lens = chunk_ends - chunk_starts
1504
+ chunk_seq_lens.clamp_(min=0)
1505
1506
cu_seq_lens_cpu = torch.zeros(
1507
num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
@@ -1550,9 +1551,8 @@ def build_mla_chunked_context_metadata(
1551
padded_local_context_lens_cpu.unsqueeze(0),
1552
local_chunk_starts + padded_local_max_context_chunk,
1553
)
1553
- padded_local_chunk_seq_lens = (local_chunk_ends - local_chunk_starts).clamp(
1554
- min=0
1555
- )
1554
+ padded_local_chunk_seq_lens = local_chunk_ends - local_chunk_starts
1555
+ padded_local_chunk_seq_lens.clamp_(min=0)
1556
1557
padded_local_cu_seq_lens_cpu = torch.zeros(
1558
num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
vllm/models/deepseek_v4/sparse_mla.py
+1
-1
@@ -209,7 +209,7 @@ class DeepseekV4FlashMLAMetadataBuilder(
209
cm.num_actual_tokens,
210
cm.query_start_loc,
211
cm.seq_lens,
212
- cm.block_table_tensor.clamp(min=0),
212
+ cm.block_table_tensor.clamp_(min=0),
213
int(self.kv_cache_spec.storage_block_size),
214
self.compress_ratio,
215
out=self.compressed_slot_mapping_buffer,
vllm/v1/attention/backends/mamba_attn.py
+24
-12
@@ -341,23 +341,36 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
341
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
342
num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
343
# Block index of the last computed token
344
- block_idx_last_computed_token = cdiv(num_computed_tokens, mamba_block_size) - 1
344
+ block_idx_last_computed_token = (
345
+ torch.div(
346
+ num_computed_tokens + mamba_block_size - 1,
347
+ mamba_block_size,
348
+ rounding_mode="floor",
349
+ )
350
+ - 1
351
+ )
352
# which is <= block index for the first scheduled token
353
block_idx_first_scheduled_token = (
347
- cdiv(num_computed_tokens + 1, mamba_block_size) - 1
354
+ torch.div(
355
+ num_computed_tokens + mamba_block_size,
356
+ mamba_block_size,
357
+ rounding_mode="floor",
358
+ )
359
+ - 1
360
)
361
# which is <= block index of the last scheduled token
362
block_idx_last_scheduled_token = (
351
- cdiv(common_attn_metadata.seq_lens, mamba_block_size) - 1
363
+ torch.div(
364
+ common_attn_metadata.seq_lens + mamba_block_size - 1,
365
+ mamba_block_size,
366
+ rounding_mode="floor",
367
+ )
368
+ - 1
369
)
370
# -1 in case it's non-computed and causes later issues with indexing
354
- block_idx_last_computed_token = torch.clamp(
355
- block_idx_last_computed_token, min=0
356
- )
371
+ block_idx_last_computed_token.clamp_(min=0)
372
# -1 in the case we have a padded request (0 seq-len)
358
- block_idx_last_scheduled_token = torch.clamp(
359
- block_idx_last_scheduled_token, min=0
360
- )
373
+ block_idx_last_scheduled_token.clamp_(min=0)
374
375
return (
376
block_idx_last_computed_token,
@@ -447,9 +460,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
460
common_attn_metadata, mamba_block_size
461
)
462
if self.use_spec_decode and prev_last_scheduled_idx is not None:
450
- fallback = torch.clamp(
451
- (num_computed_tokens - 1) // mamba_block_size, min=0
452
- )
463
+ fallback = (num_computed_tokens - 1) // mamba_block_size
464
+ fallback.clamp_(min=0)
465
block_idx_last_scheduled_token_prev_step = torch.where(
466
prev_last_scheduled_idx >= 0,
467
prev_last_scheduled_idx,
vllm/v1/attention/backends/rocm_aiter_fa.py
+2
-3
@@ -593,9 +593,8 @@ class AiterFlashAttentionMetadataBuilder(
593
chunk_ends = torch.min(
594
computed_kv_lens.unsqueeze(0), chunk_starts + max_context_chunk
595
)
596
- chunk_seq_lens = (chunk_ends - chunk_starts).clamp(
597
- min=0
598
- ) # [num_chunks, num_extends]
596
+ chunk_seq_lens = chunk_ends - chunk_starts
597
+ chunk_seq_lens.clamp_(min=0) # [num_chunks, num_extends]
598
cu_seq_lens_cpu = torch.zeros(
599
[num_chunks, num_extends + 1], dtype=torch.int32, pin_memory=True
600
)
vllm/v1/attention/backends/utils.py
+2
-4
@@ -950,10 +950,8 @@ def mamba_get_block_table_tensor(
950
assert isinstance(kv_cache_spec, MambaSpec)
951
# NOTE: For 0-length requests in CUDA graph, use a start_index of 0
952
# to handle the invalid block table.
953
- start_indices = torch.clamp(
954
- (seq_lens - 1) // kv_cache_spec.block_size,
955
- min=0,
956
- )
953
+ start_indices = (seq_lens - 1) // kv_cache_spec.block_size
954
+ start_indices.clamp_(min=0)
955
# Use int32 for arithmetic to avoid dtype promotion overhead,
956
# then convert to int64 for gather (which requires Long indices)
957
offsets = torch.arange(
vllm/v1/spec_decode/ngram_proposer_gpu.py
+12
-7
@@ -91,7 +91,8 @@ class NgramGPUKernel(nn.Module):
91
suffix_indices = suffix_starts.unsqueeze(1) + torch.arange(
92
ngram_len, device=device
93
)
94
- suffix = torch.gather(token_ids, 1, suffix_indices.clamp(min=0))
94
+ suffix_indices.clamp_(min=0)
95
+ suffix = torch.gather(token_ids, 1, suffix_indices)
96
97
# Window matches for each sequence.
98
matches = (search_windows == suffix.unsqueeze(1)).all(dim=-1)
@@ -134,7 +135,7 @@ class NgramGPUKernel(nn.Module):
135
draft_indices = draft_start.unsqueeze(1) + torch.arange(
136
num_draft_tokens, device=device
137
)
137
- draft_indices = draft_indices.clamp(min=0, max=max_seq_len - 1)
138
+ draft_indices.clamp_(min=0, max=max_seq_len - 1)
139
140
# Extract draft tokens; gather always runs.
141
draft_tokens = torch.gather(token_ids, 1, draft_indices)
@@ -357,7 +358,8 @@ class NgramProposerGPU:
358
valid_write_mask & (valid_sampled_token_ids_gpu != -1) & in_bounds
359
)
360
360
- write_positions_long = write_positions.clamp(max=max_seq_len - 1).long()
361
+ write_positions.clamp_(max=max_seq_len - 1)
362
+ write_positions_long = write_positions.long()
363
existing_values = token_ids_gpu.gather(1, write_positions_long)
364
365
tokens_cast = valid_sampled_token_ids_gpu.to(token_ids_gpu.dtype)
@@ -427,7 +429,9 @@ class NgramProposerGPU:
429
)
430
431
# Backup last valid token before speculative tokens.
430
- backup_indices = (num_tokens_no_spec[:num_reqs] - 1).clamp(min=0).long()
432
+ backup_indices = num_tokens_no_spec[:num_reqs] - 1
433
+ backup_indices.clamp_(min=0)
434
+ backup_indices = backup_indices.long()
435
backup_next_token_ids = torch.gather(
436
token_ids_gpu[:num_reqs], dim=1, index=backup_indices.unsqueeze(1)
437
).squeeze(1)
@@ -447,16 +451,17 @@ class NgramProposerGPU:
451
452
# Rightmost valid index per row.
453
last_valid_indices = valid_sampled_tokens_count - 1
450
- last_valid_indices_safe = torch.clamp(last_valid_indices, min=0)
454
+ has_valid_sample = last_valid_indices >= 0
455
+ last_valid_indices.clamp_(min=0)
456
457
# Last valid token from each row; undefined if none.
458
selected_tokens = torch.gather(
454
- valid_sampled_token_ids_gpu, 1, last_valid_indices_safe.unsqueeze(1)
459
+ valid_sampled_token_ids_gpu, 1, last_valid_indices.unsqueeze(1)
460
).squeeze(1)
461
462
# Use last token if valid; otherwise fallback to backup.
463
next_token_ids = torch.where(
459
- last_valid_indices != -1,
464
+ has_valid_sample,
465
selected_tokens,
466
backup_next_token_ids,
467
)
vllm/v1/spec_decode/step3p5.py
+3
-1
@@ -107,7 +107,9 @@ class Step3p5MTPProposer(EagleProposer):
107
if block_table is None:
108
continue
109
n_blocks = block_table.shape[1]
110
- bn = (new_positions_1d // block_size).clamp(max=n_blocks - 1).to(torch.long)
110
+ bn = new_positions_1d // block_size
111
+ bn.clamp_(max=n_blocks - 1)
112
+ bn = bn.to(torch.long)
113
block_ids = block_table[:batch_size].gather(1, bn.unsqueeze(1)).squeeze(1)
114
sm = block_ids * block_size + (new_positions_1d % block_size)
115
sm.masked_fill_(exceeds, PADDING_SLOT_ID)