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 6, 13:15 UTC
[ROCm][Quantization] add fp8xfp8 attn support for rocm_aiter_unified_attn (#36927)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Mar 17, 2026 at 20:49 UTC
e6c4797704e45c4b157966c197a1cc53d0b92237
1 file changed
+17
-5
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+17
-5
@@ -125,6 +125,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
125
from aiter.ops.triton.unified_attention import unified_attention
126
127
self.unified_attention = unified_attention
128
+ self.supports_quant_query_input = True
129
130
def forward(
131
self,
@@ -190,12 +191,20 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
191
192
key_cache, value_cache = kv_cache.unbind(0)
193
194
+ softmax_scale = self.scale
195
+ fp8_post_attn_v_rescale = False
196
if self.kv_cache_dtype.startswith("fp8"):
197
key_cache = key_cache.view(self.fp8_dtype)
198
value_cache = value_cache.view(self.fp8_dtype)
196
- assert layer._q_scale_float == 1.0, (
197
- "A non 1.0 q_scale is not currently supported."
198
- )
199
+ # When Q is FP8, triton kernel skips K/V dequant (for fp8xfp8 matmul).
200
+ # Compensate by absorbing q_scale and k_scale into softmax_scale, and
201
+ # v_scale into output_scale (or post-multiplying if no fusion).
202
+ if query.dtype == self.fp8_dtype:
203
+ softmax_scale = self.scale * layer._q_scale_float * layer._k_scale_float
204
+ if output_scale is not None:
205
+ output_scale = output_scale / layer._v_scale_float
206
+ else:
207
+ fp8_post_attn_v_rescale = True
208
209
cu_seqlens_q = attn_metadata.query_start_loc
210
seqused_k = attn_metadata.seq_lens
@@ -217,19 +226,22 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
226
max_seqlen_q=max_seqlen_q,
227
seqused_k=seqused_k,
228
max_seqlen_k=max_seqlen_k,
220
- softmax_scale=self.scale,
229
+ softmax_scale=softmax_scale,
230
causal=True,
231
alibi_slopes=self.alibi_slopes,
232
window_size=self.sliding_window,
233
block_table=block_table,
234
softcap=self.logits_soft_cap,
226
- q_descale=None, # Not supported
235
+ q_descale=None, # q_scale absorbed into softmax_scale
236
k_descale=layer._k_scale.expand(descale_shape),
237
v_descale=layer._v_scale.expand(descale_shape),
238
sinks=self.sinks,
239
output_scale=output_scale,
240
)
241
242
+ if fp8_post_attn_v_rescale:
243
+ output[:num_actual_tokens].mul_(layer._v_scale_float)
244
+
245
return output
246
247
def do_kv_cache_update(