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 7, 14:00 UTC
[ROCm] Enabling encoder and encoder-decoder on ROCm and AITER unified backends (#35334)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Feb 27, 2026 at 15:32 UTC
9fa6c68fa627c7ab041c48ac9987fb093719597f
3 files changed
+106
-7
docs/design/attention_backends.md
+2
-2
@@ -171,8 +171,8 @@ Priority is **1 = highest** (tried first).
171
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
172
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
173
| `ROCM_AITER_FA` | | fp16, bf16 | `auto` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
174
-| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | Decoder | N/A |
175
-| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
174
+| `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | All | N/A |
175
+| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto` | 16, 32, 544 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A |
176
| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any |
177
| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any |
178
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+31
@@ -55,6 +55,16 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
55
def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]:
56
return RocmAttentionMetadataBuilder
57
58
+ @classmethod
59
+ def supports_attn_type(cls, attn_type: str) -> bool:
60
+ """RocmAiterUnifiedAttention supports all attention types."""
61
+ return attn_type in (
62
+ AttentionType.DECODER,
63
+ AttentionType.ENCODER,
64
+ AttentionType.ENCODER_ONLY,
65
+ AttentionType.ENCODER_DECODER,
66
+ )
67
+
68
69
class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
70
def fused_output_quant_supported(self, quant_key: QuantKey):
@@ -143,6 +153,19 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
153
154
num_actual_tokens = attn_metadata.num_actual_tokens
155
156
+ # Handle encoder attention differently - no KV cache needed
157
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
158
+ # For encoder attention,
159
+ # we use direct Q, K, V tensors without caching
160
+ return self._forward_encoder_attention(
161
+ query[:num_actual_tokens],
162
+ key[:num_actual_tokens],
163
+ value[:num_actual_tokens],
164
+ output[:num_actual_tokens],
165
+ attn_metadata,
166
+ layer,
167
+ )
168
+
169
key_cache, value_cache = kv_cache.unbind(0)
170
171
if self.kv_cache_dtype.startswith("fp8"):
@@ -195,6 +218,10 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
218
kv_cache: torch.Tensor,
219
slot_mapping: torch.Tensor,
220
):
221
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
222
+ # For encoder attention,
223
+ # we use direct Q, K, V tensors without caching
224
+ return
225
key_cache, value_cache = kv_cache.unbind(0)
226
227
# Reshape the input keys and values and store them in the cache.
@@ -224,6 +251,10 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
251
kv_cache: torch.Tensor,
252
layer_slot_mapping: torch.Tensor,
253
):
254
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
255
+ # For encoder attention,
256
+ # we use direct Q, K, V tensors without caching
257
+ return
258
key_cache, value_cache = kv_cache.unbind(0)
259
flash_layout = True
260
vllm/v1/attention/backends/rocm_attn.py
+73
-5
@@ -205,6 +205,16 @@ class RocmAttentionBackend(AttentionBackend):
205
def get_impl_cls() -> type["RocmAttentionImpl"]:
206
return RocmAttentionImpl
207
208
+ @classmethod
209
+ def supports_attn_type(cls, attn_type: str) -> bool:
210
+ """RocmAttention supports all attention types."""
211
+ return attn_type in (
212
+ AttentionType.DECODER,
213
+ AttentionType.ENCODER,
214
+ AttentionType.ENCODER_ONLY,
215
+ AttentionType.ENCODER_DECODER,
216
+ )
217
+
218
@staticmethod
219
def get_kv_cache_shape(
220
num_blocks: int,
@@ -244,6 +254,7 @@ class RocmAttentionImpl(AttentionImpl):
254
kv_sharing_target_layer_name: int | None = None,
255
sinks: torch.Tensor | None = None,
256
) -> None:
257
+ self.attn_type = attn_type
258
self.num_heads = num_heads
259
self.head_size = head_size
260
self.scale = float(scale)
@@ -266,11 +277,6 @@ class RocmAttentionImpl(AttentionImpl):
277
278
RocmAttentionBackend.validate_head_size(head_size)
279
269
- if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
270
- raise NotImplementedError(
271
- "Encoder self-attention is not implemented for RocmAttentionImpl"
272
- )
273
-
280
self.fp8_dtype = current_platform.fp8_dtype()
281
282
self.sinks = sinks
@@ -281,6 +287,54 @@ class RocmAttentionImpl(AttentionImpl):
287
f"num_heads: {num_heads}."
288
)
289
290
+ def _forward_encoder_attention(
291
+ self,
292
+ query: torch.Tensor,
293
+ key: torch.Tensor,
294
+ value: torch.Tensor,
295
+ output: torch.Tensor,
296
+ attn_metadata: FlashAttentionMetadata,
297
+ layer: torch.nn.Module,
298
+ ) -> torch.Tensor:
299
+ """Forward pass for encoder attention without KV cache.
300
+
301
+ Args:
302
+ query: shape = [num_encoder_tokens, num_heads, head_size]
303
+ key: shape = [num_encoder_tokens, num_kv_heads, head_size]
304
+ value: shape = [num_encoder_tokens, num_kv_heads, head_size]
305
+ output: shape = [num_encoder_tokens, num_heads, head_size]
306
+ attn_metadata: Encoder attention metadata
307
+ layer: The attention layer
308
+ """
309
+ # For encoder attention, process FP8 quantization if needed
310
+ if self.kv_cache_dtype.startswith("fp8"):
311
+ raise NotImplementedError(
312
+ "quantization is not supported for encoder attention"
313
+ )
314
+
315
+ # Use encoder-specific metadata for sequence information
316
+ query_start_loc = attn_metadata.query_start_loc
317
+ seq_lens = attn_metadata.seq_lens
318
+ max_query_len = attn_metadata.max_query_len
319
+
320
+ # Call flash attention directly on Q, K, V tensors
321
+ from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
322
+
323
+ context_attention_fwd(
324
+ q=query,
325
+ k=key,
326
+ v=value,
327
+ o=output,
328
+ b_start_loc=query_start_loc,
329
+ b_seq_len=seq_lens,
330
+ max_input_len=max_query_len,
331
+ is_causal=False,
332
+ softmax_scale=self.scale,
333
+ sliding_window_q=self.sliding_window[0],
334
+ sliding_window_k=self.sliding_window[1],
335
+ )
336
+ return output
337
+
338
def forward(
339
self,
340
layer: torch.nn.Module,
@@ -330,6 +384,16 @@ class RocmAttentionImpl(AttentionImpl):
384
385
num_actual_tokens = attn_metadata.num_actual_tokens
386
387
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
388
+ return self._forward_encoder_attention(
389
+ query[:num_actual_tokens],
390
+ key[:num_actual_tokens],
391
+ value[:num_actual_tokens],
392
+ output[:num_actual_tokens],
393
+ attn_metadata,
394
+ layer,
395
+ )
396
+
397
key_cache, value_cache = PagedAttention.split_kv_cache(
398
kv_cache, self.num_kv_heads, self.head_size
399
)
@@ -380,6 +444,8 @@ class RocmAttentionImpl(AttentionImpl):
444
kv_cache: torch.Tensor,
445
slot_mapping: torch.Tensor,
446
):
447
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
448
+ return
449
key_cache, value_cache = PagedAttention.split_kv_cache(
450
kv_cache, self.num_kv_heads, self.head_size
451
)
@@ -432,6 +498,8 @@ class RocmAttentionImpl(AttentionImpl):
498
kv_cache: torch.Tensor,
499
layer_slot_mapping: torch.Tensor,
500
):
501
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
502
+ return
503
key_cache, value_cache = PagedAttention.split_kv_cache(
504
kv_cache,
505
layer.num_kv_heads, # type: ignore[attr-defined]