@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6d09769700

Fork 0

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 4, 20:00 UTC

[ROCm] Support non-causal attention in ROCM_ATTN (#40176)

Signed-off-by: Micah Williamson <micah.williamson@amd.com>

Micah Williamson committed Apr 21, 2026 at 22:57 UTC 6d097697001ed8cf06190076b70ce36c8bab3437

4 files changed +41 -13

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+6 -2

@@ -13,10 +13,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

13 )

14 from vllm.utils.torch_utils import is_quantized_kv_cache

15 from vllm.v1.attention.backend import AttentionLayer, AttentionType, MultipleOf

16 -from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

16 from vllm.v1.attention.backends.rocm_attn import (

17 RocmAttentionBackend,

18 RocmAttentionImpl,

19 + RocmAttentionMetadata,

20 RocmAttentionMetadataBuilder,

21 )

@@ -53,6 +53,10 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):

53 def supports_sink(cls) -> bool:

54 return True

56 + @classmethod

57 + def supports_non_causal(cls) -> bool:

58 + return False

59 +

60 forward_includes_kv_cache_update: bool = False

62 @staticmethod

@@ -140,7 +144,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

144 key: torch.Tensor,

145 value: torch.Tensor,

146 kv_cache: torch.Tensor,

143 - attn_metadata: FlashAttentionMetadata,

147 + attn_metadata: RocmAttentionMetadata,

148 output: torch.Tensor,

149 output_scale: torch.Tensor | None = None,

150 output_block_scale: torch.Tensor | None = None,

vllm/v1/attention/backends/rocm_attn.py

+11 -3

@@ -27,7 +27,6 @@ from vllm.v1.attention.backend import (

27 CommonAttentionMetadata,

28 MultipleOf,

29 )

30 -from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata

30 from vllm.v1.attention.ops.chunked_prefill_paged_decode import (

31 chunked_prefill_paged_decode,

32 )

@@ -69,6 +68,9 @@ class RocmAttentionMetadata:

68 scheduler_metadata: torch.Tensor | None = None

69 prefix_scheduler_metadata: torch.Tensor | None = None

71 + # DFlash drafting sets this to False via CommonAttentionMetadata.

72 + causal: bool = True

73 +

75 class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadata]):

76 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS

@@ -154,6 +156,7 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat

156 prefix_kv_lens=prefix_kv_lens,

157 suffix_kv_lens=suffix_kv_lens,

158 prefix_scheduler_metadata=prefix_scheduler_metadata,

159 + causal=common_attn_metadata.causal,

160 )

161 return attn_metadata

162

@@ -200,6 +203,10 @@ class RocmAttentionBackend(AttentionBackend):

203 # kernel, which is less efficient than the proper triton backends.

204 return False

205

206 + @classmethod

207 + def supports_non_causal(cls) -> bool:

208 + return True

209 +

210 forward_includes_kv_cache_update: bool = False

211

212 @staticmethod

@@ -301,7 +308,7 @@ class RocmAttentionImpl(AttentionImpl):

308 key: torch.Tensor,

309 value: torch.Tensor,

310 output: torch.Tensor,

304 - attn_metadata: FlashAttentionMetadata,

311 + attn_metadata: RocmAttentionMetadata,

312 layer: torch.nn.Module,

313 ) -> torch.Tensor:

314 """Forward pass for encoder attention without KV cache.

@@ -350,7 +357,7 @@ class RocmAttentionImpl(AttentionImpl):

357 key: torch.Tensor,

358 value: torch.Tensor,

359 kv_cache: torch.Tensor,

353 - attn_metadata: FlashAttentionMetadata,

360 + attn_metadata: RocmAttentionMetadata,

361 output: torch.Tensor,

362 output_scale: torch.Tensor | None = None,

363 output_block_scale: torch.Tensor | None = None,

@@ -438,6 +445,7 @@ class RocmAttentionImpl(AttentionImpl):

445 sm_scale=self.scale,

446 output_scale=output_scale,

447 sinks=self.sinks,

448 + causal=attn_metadata.causal,

449 )

450

451 return output

vllm/v1/attention/ops/chunked_prefill_paged_decode.py

@@ -269,6 +269,7 @@ def chunked_prefill_paged_decode(

269 # Optional tensor for sinks

270 sinks=None,

271 is_block_table_ptr: bool = False,

272 + causal: bool = True,

273 ):

274 if sm_scale is None:

275 sm_scale = 1.0 / (query.shape[2] ** 0.5)

@@ -300,6 +301,7 @@ def chunked_prefill_paged_decode(

301 skip_decode=True,

302 fp8_out_scale=output_scale,

303 sinks=sinks,

304 + causal=causal,

305 )

306

307 block_size = value_cache.shape[3]

vllm/v1/attention/ops/prefix_prefill.py

+22 -8

@@ -89,6 +89,7 @@ def _fwd_kernel(

89 SKIP_DECODE: tl.constexpr,

90 USE_SINKS: tl.constexpr,

91 USE_FP8: tl.constexpr,

92 + CAUSAL: tl.constexpr = True,

93 MAX_Q_LEN: tl.constexpr = 0,

94 MAX_CTX_LEN: tl.constexpr = 0,

95 FP8_MIN: tl.constexpr = float8_info.min,

@@ -283,10 +284,17 @@ def _fwd_kernel(

284 # block_mask is 0 when we're already past the current query length

285 block_mask = tl.where(block_start_loc < cur_batch_query_len, 1, 0)

286

286 - # compute query against itself (with causal mask)

287 + # compute query against itself (causal among queries by default;

288 + # CAUSAL=False for bidirectional attention over query tokens, e.g. DFlash.)

289 + if CAUSAL:

290 + key_range_upper = block_mask * (start_m + 1) * BLOCK_M

291 + else:

292 + q_len_pad = (cur_batch_query_len + BLOCK_N - 1) // BLOCK_N * BLOCK_N

293 + key_range_upper = block_mask * q_len_pad

294 +

295 for start_n in tl.range(

296 0,

289 - block_mask * (start_m + 1) * BLOCK_M,

297 + key_range_upper,

298 BLOCK_N,

299 loop_unroll_factor=num_unroll_request,

300 ):

@@ -302,14 +310,17 @@ def _fwd_kernel(

310 qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)

311 qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)

312 qk *= sm_scale

305 - # apply causal mask

306 - qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))

313 +

314 + valid_kv = (start_n + offs_n[None, :]) < cur_batch_query_len

315 + if CAUSAL:

316 + attn_mask = valid_kv & (offs_m[:, None] >= (start_n + offs_n[None, :]))

317 + else:

318 + attn_mask = valid_kv

319 if SLIDING_WINDOW > 0:

308 - qk = tl.where(

309 - offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW,

310 - qk,

311 - float("-inf"),

320 + attn_mask = attn_mask & (

321 + offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW

322 )

323 + qk = tl.where(attn_mask, qk, float("-inf"))

324

325 # compute running maximum

326 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))

@@ -656,6 +667,7 @@ def context_attention_fwd(

667 fp8_out_scale=None,

668 sinks=None,

669 is_block_table_ptr: bool = False,

670 + causal: bool = True,

671 ):

672 q_dtype_is_f32 = q.dtype is torch.float32

673

@@ -722,6 +734,7 @@ def context_attention_fwd(

734 processed_b_loc = b_loc.to(torch.int32)

735

736 if alibi_slopes is not None:

737 + assert causal, "Non-causal prefix attention is not supported with alibi"

738 assert sinks is None, "Sinks arg is not supported with alibi"

739 assert fp8_out_scale is None, "FP8 output not supported with alibi"

740 # need to reduce num. blocks when using fp32

@@ -859,6 +872,7 @@ def context_attention_fwd(

872 num_warps=4,

873 num_stages=1,

874 USE_SINKS=sinks is not None,

875 + CAUSAL=causal,

876 **extra_kargs,

877 )

878 return