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, 05:00 UTC
[ROCm] aiter_unified_attn fp8 q scale refactor (#38296)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
May 6, 2026 at 12:11 UTC
22a3cbe1520bc8a3b19ace0abe497c514b3129ea
1 file changed
+3
-21
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+3
-21
@@ -200,19 +200,9 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
200
key_cache, value_cache = kv_cache.unbind(0)
201
202
softmax_scale = self.scale
203
- fp8_post_attn_v_rescale = False
203
if is_quantized_kv_cache(self.kv_cache_dtype):
204
key_cache = key_cache.view(self.fp8_dtype)
205
value_cache = value_cache.view(self.fp8_dtype)
207
- # When Q is FP8, triton kernel skips K/V dequant (for fp8xfp8 matmul).
208
- # Compensate by absorbing q_scale and k_scale into softmax_scale, and
209
- # v_scale into output_scale (or post-multiplying if no fusion).
210
- if query.dtype == self.fp8_dtype:
211
- softmax_scale = self.scale * layer._q_scale_float * layer._k_scale_float
212
- if output_scale is not None:
213
- output_scale = output_scale / layer._v_scale_float
214
- else:
215
- fp8_post_attn_v_rescale = True
206
207
cu_seqlens_q = attn_metadata.query_start_loc
208
seqused_k = attn_metadata.seq_lens
@@ -220,11 +210,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
210
max_seqlen_k = attn_metadata.max_seq_len
211
block_table = attn_metadata.block_table
212
223
- descale_shape = (
224
- cu_seqlens_q.shape[0] - 1,
225
- key.shape[1] if key is not None else self.num_kv_heads,
226
- )
227
-
213
self.unified_attention(
214
q=query[:num_actual_tokens],
215
k=key_cache,
@@ -240,16 +225,13 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
225
window_size=self.sliding_window,
226
block_table=block_table,
227
softcap=self.logits_soft_cap,
243
- q_descale=None, # q_scale absorbed into softmax_scale
244
- k_descale=layer._k_scale.expand(descale_shape),
245
- v_descale=layer._v_scale.expand(descale_shape),
228
+ q_descale=layer._q_scale if query.dtype == self.fp8_dtype else None,
229
+ k_descale=layer._k_scale,
230
+ v_descale=layer._v_scale,
231
sinks=self.sinks,
232
output_scale=output_scale,
233
)
234
250
- if fp8_post_attn_v_rescale:
251
- output[:num_actual_tokens].mul_(layer._v_scale_float)
252
-
235
return output
236
237
def do_kv_cache_update(