Update test_flash_attn.py (#17102)

Signed-off-by: ShuaibinLi <lishuaibin@live.cn>
This commit is contained in:
Happy 2025-04-27 06:17:35 +08:00 committed by GitHub
parent 3642c59aa8
commit 9869453c42
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -145,7 +145,7 @@ def test_flash_attn_with_paged_kv(
v_descale = None
if q_dtype is not None:
# QKV are drawn from N(0, 1): no need for a fp8 scaling factor
maybe_quantized_query = query.to(q_dtype)
maybe_quantized_query = q.to(q_dtype)
maybe_quantized_key_cache = key_cache.to(q_dtype)
maybe_quantized_value_cache = value_cache.to(q_dtype)