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 15, 01:00 UTC
[Attention] Clean up iRoPE in V1 (#21188)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Lucas Wilkinson committed
Jul 21, 2025 at 12:10 UTC
304dce7ec02769ecea137091caa5413e1a4abf60
9 files changed
+14
-26
vllm/attention/layer.py
+7
@@ -137,6 +137,13 @@ class Attention(nn.Module):
137
self.num_kv_heads = num_kv_heads
138
self.sliding_window = sliding_window
139
140
+ # For v1 we have backend agnostic iRoPE (local chunked attention)
141
+ # we have to store the flag on the layer so gpu model runner can
142
+ # set KVSpec appropriately (and pop it so it doesnt get passed to
143
+ # the backends)
144
+ if envs.VLLM_USE_V1:
145
+ self.use_irope = extra_impl_args.pop("use_irope", False)
146
+
147
quant_method = quant_config.get_quant_method(
148
self, prefix=prefix) if quant_config else None
149
if quant_method is not None and not isinstance(
vllm/v1/attention/backends/cpu_attn.py
-5
@@ -446,17 +446,12 @@ class TorchSDPABackendImpl(AttentionImpl[TorchSDPAMetadata]):
446
logits_soft_cap: Optional[float] = None,
447
attn_type: str = AttentionType.DECODER,
448
kv_sharing_target_layer_name: Optional[str] = None,
449
- use_irope: bool = False,
449
) -> None:
450
if kv_sharing_target_layer_name is not None:
451
raise NotImplementedError("KV sharing is not supported in V0.")
452
if logits_soft_cap is not None:
453
logger.warning_once("Torch SPDA does not support logits soft cap. "
454
"Outputs may be slightly off.")
456
- if use_irope:
457
- logger.warning_once(
458
- "Using irope in Torch SPDA is not supported yet, it will fall"
459
- " back to global attention for long context.")
455
self.paged_attn_impl = _get_paged_attn_impl()
456
self.num_heads = num_heads
457
self.head_size = head_size
vllm/v1/attention/backends/flash_attn.py
-2
@@ -352,7 +352,6 @@ class FlashAttentionImpl(AttentionImpl):
352
logits_soft_cap: Optional[float] = None,
353
attn_type: AttentionType = AttentionType.DECODER,
354
kv_sharing_target_layer_name: Optional[str] = None,
355
- use_irope: bool = False,
355
) -> None:
356
self.num_heads = num_heads
357
self.head_size = head_size
@@ -381,7 +380,6 @@ class FlashAttentionImpl(AttentionImpl):
380
"encoder/decoder cross-attention "
381
"are not implemented for "
382
"FlashAttentionImpl")
384
- self.use_irope = use_irope
383
self.vllm_flash_attn_version = get_flash_attn_version()
384
if is_quantized_kv_cache(self.kv_cache_dtype) \
385
and not flash_attn_supports_fp8():
vllm/v1/attention/backends/flashinfer.py
-2
@@ -493,7 +493,6 @@ class FlashInferImpl(AttentionImpl):
493
logits_soft_cap: Optional[float] = None,
494
attn_type: AttentionType = AttentionType.DECODER,
495
kv_sharing_target_layer_name: Optional[int] = None,
496
- use_irope: bool = False,
496
) -> None:
497
self.num_heads = num_heads
498
self.head_size = head_size
@@ -509,7 +508,6 @@ class FlashInferImpl(AttentionImpl):
508
self.kv_cache_dtype = kv_cache_dtype
509
self.logits_soft_cap = logits_soft_cap
510
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
512
- self.use_irope = use_irope
511
512
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
513
vllm/v1/attention/backends/pallas.py
-5
@@ -148,12 +148,7 @@ class PallasAttentionBackendImpl(AttentionImpl):
148
logits_soft_cap: Optional[float] = None,
149
attn_type: str = AttentionType.DECODER,
150
kv_sharing_target_layer_name: Optional[int] = None,
151
- use_irope: bool = False,
151
) -> None:
153
- if use_irope:
154
- logger.warning_once(
155
- "Using irope in Pallas is not supported yet, it will fall back "
156
- "to global attention for long context.")
152
self.num_heads = num_heads
153
self.head_size = head_size
154
self.scale = float(scale)
vllm/v1/attention/backends/rocm_aiter_fa.py
-2
@@ -337,7 +337,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
337
logits_soft_cap: Optional[float] = None,
338
attn_type: AttentionType = AttentionType.DECODER,
339
kv_sharing_target_layer_name: Optional[int] = None,
340
- use_irope: bool = False,
340
) -> None:
341
self.num_heads = num_heads
342
self.head_size = head_size
@@ -367,7 +366,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
366
"encoder/decoder cross-attention "
367
"are not implemented for "
368
"FlashAttentionImpl")
370
- self.use_irope = use_irope
369
if is_quantized_kv_cache(self.kv_cache_dtype):
370
raise NotImplementedError(
371
"AiterFlashAttention does not support fp8 kv-cache on this "
vllm/v1/attention/backends/triton_attn.py
-6
@@ -72,9 +72,6 @@ class TritonAttentionMetadataBuilder(
72
vllm_config.parallel_config)
73
self.headdim = model_config.get_head_size()
74
75
- self.attention_chunk_size = getattr(vllm_config.scheduler_config,
76
- 'attention_chunk_size', None)
77
-
75
def build_for_cudagraph_capture(
76
self, common_attn_metadata: CommonAttentionMetadata
77
) -> TritonAttentionMetadata:
@@ -208,7 +205,6 @@ class TritonAttentionImpl(AttentionImpl):
205
logits_soft_cap: Optional[float] = None,
206
attn_type: AttentionType = AttentionType.DECODER,
207
kv_sharing_target_layer_name: Optional[int] = None,
211
- use_irope: bool = False,
208
) -> None:
209
self.num_heads = num_heads
210
self.head_size = head_size
@@ -228,8 +224,6 @@ class TritonAttentionImpl(AttentionImpl):
224
self.logits_soft_cap = logits_soft_cap
225
self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
226
231
- self.use_irope = use_irope
232
-
227
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
228
229
TritonAttentionBackend.validate_head_size(head_size)
vllm/v1/worker/gpu_model_runner.py
+3
-4
@@ -2702,8 +2702,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2702
# TODO: Support other attention modules, e.g., cross-attention
2703
if attn_module.attn_type == AttentionType.DECODER:
2704
use_local_attention = (self.attention_chunk_size is not None
2705
- and getattr(attn_module.impl,
2706
- "use_irope", False))
2705
+ and attn_module.use_irope)
2706
if attn_module.sliding_window is not None:
2707
kv_cache_spec[layer_name] = SlidingWindowSpec(
2708
block_size=block_size,
@@ -2716,13 +2715,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2715
"attention module can not be with ",
2716
"both local attention and sliding window")
2717
elif use_local_attention:
2719
- kv_cache_spec[layer_name] = (ChunkedLocalAttentionSpec(
2718
+ kv_cache_spec[layer_name] = ChunkedLocalAttentionSpec(
2719
block_size=block_size,
2720
num_kv_heads=attn_module.num_kv_heads,
2721
head_size=attn_module.head_size,
2722
dtype=self.kv_cache_dtype,
2723
attention_chunk_size=self.attention_chunk_size,
2725
- use_mla=use_mla))
2724
+ use_mla=use_mla)
2725
else:
2726
kv_cache_spec[layer_name] = FullAttentionSpec(
2727
block_size=block_size,
vllm/v1/worker/tpu_model_runner.py
+4
@@ -519,6 +519,10 @@ class TPUModelRunner(LoRAModelRunnerMixin):
519
continue
520
521
if attn_module.attn_type == AttentionType.DECODER:
522
+ if attn_module.use_irope:
523
+ logger.warning_once(
524
+ "Using irope in Pallas is not supported yet, it "
525
+ "will fall back to global attention for long context.")
526
if attn_module.sliding_window is not None:
527
kv_cache_spec[layer_name] = SlidingWindowSpec(
528
block_size=block_size,