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, 11:15 UTC
[Bugfix]: Fix ROCm fusion attn test; use AttentionBackend utils to create kv cache (#33948)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Feb 11, 2026 at 10:12 UTC
fd618871b41c0cf9259379cde9cca230a56c4096
1 file changed
+27
-52
tests/compile/passes/test_fusion_attn.py
+27
-52
@@ -92,6 +92,8 @@ class AttentionQuantPatternModel(torch.nn.Module):
92
def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
93
"""Initialize attention metadata."""
94
95
+ # TODO (Rohan138) reuse utils from vllm/v1/worker/gpu/attn_utils.py
96
+
97
# Create common attn metadata
98
batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
99
common_attn_metadata = create_common_attn_metadata(
@@ -100,58 +102,31 @@ class AttentionQuantPatternModel(torch.nn.Module):
102
103
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
104
num_blocks = batch_size * max_blocks
103
- backend = self.attn.backend
104
-
105
- # TODO(luka) use get_kv_cache_stride_order
106
- # Create dummy KV cache for the selected backend
107
- if backend == AttentionBackendEnum.ROCM_ATTN:
108
- # k/v as 1st dimention
109
- # HND: [num_blocks, num_kv_heads, block_size, head_size]
110
- kv_cache = torch.zeros(
111
- 2,
112
- num_blocks,
113
- self.num_kv_heads,
114
- self.block_size,
115
- self.head_size,
116
- dtype=self.kv_cache_dtype,
117
- device=self.device,
118
- )
119
- elif backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN:
120
- # k/v as 1st dimention
121
- # NHD: [num_blocks, block_size, num_kv_heads, head_size]
122
- kv_cache = torch.zeros(
123
- 2,
124
- num_blocks,
125
- self.block_size,
126
- self.num_kv_heads,
127
- self.head_size,
128
- dtype=self.kv_cache_dtype,
129
- device=self.device,
130
- )
131
- elif backend == AttentionBackendEnum.TRITON_ATTN:
132
- # k/v as 2nd dimention
133
- # NHD: [num_blocks, block_size, num_kv_heads, head_size]
134
- kv_cache = torch.zeros(
135
- num_blocks,
136
- 2,
137
- self.num_kv_heads,
138
- self.block_size,
139
- self.head_size,
140
- dtype=self.kv_cache_dtype,
141
- device=self.device,
142
- )
143
- elif backend == AttentionBackendEnum.FLASHINFER:
144
- kv_cache = torch.zeros(
145
- num_blocks,
146
- 2,
147
- self.num_kv_heads,
148
- self.block_size,
149
- self.head_size,
150
- dtype=self.kv_cache_dtype,
151
- device=self.device,
152
- ).permute(0, 1, 3, 2, 4)
153
- else:
154
- raise ValueError(f"Unsupported backend: {backend}")
105
+
106
+ # Fetch the attention backend and kv cache shape and stride order
107
+ attn_backend = self.attn.attn_backend
108
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
109
+ num_blocks, self.block_size, self.num_kv_heads, self.head_size
110
+ )
111
+ try:
112
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
113
+ except (AttributeError, NotImplementedError):
114
+ kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
115
+
116
+ kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
117
+ inv_order = [
118
+ kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
119
+ ]
120
+
121
+ # Create dummy KV cache
122
+ raw_tensor = torch.zeros(
123
+ 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
124
+ dtype=self.kv_cache_dtype,
125
+ device=self.device,
126
+ )
127
+ raw_tensor = raw_tensor.view(kv_cache_shape)
128
+ kv_cache = raw_tensor.permute(*inv_order)
129
+
130
self.attn.kv_cache = [kv_cache]
131
132
# Build attn metadata