[AMD] Fused qk rmsnorm bf16 for amd/Kimi-K2.5-MXFP4 (#23186)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
fa8993111d
commit
8589b92a89
@@ -60,6 +60,9 @@ if _is_cuda:
|
|||||||
|
|
||||||
|
|
||||||
if _use_aiter:
|
if _use_aiter:
|
||||||
|
from aiter.ops.fused_qk_norm_rope_cache_quant import (
|
||||||
|
fused_qk_rmsnorm as fused_qk_rmsnorm_bf16,
|
||||||
|
)
|
||||||
from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
|
from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import (
|
||||||
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
|
batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant,
|
||||||
)
|
)
|
||||||
@@ -160,6 +163,15 @@ class DeepseekMLAForwardMixin:
|
|||||||
output_unquantized_inp1=False,
|
output_unquantized_inp1=False,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
elif _use_aiter:
|
||||||
|
q, k_nope = fused_qk_rmsnorm_bf16(
|
||||||
|
q,
|
||||||
|
self.q_a_layernorm.weight,
|
||||||
|
self.q_a_layernorm.variance_epsilon,
|
||||||
|
k_nope,
|
||||||
|
self.kv_a_layernorm.weight,
|
||||||
|
self.kv_a_layernorm.variance_epsilon,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
q = self.q_a_layernorm(q)
|
q = self.q_a_layernorm(q)
|
||||||
k_nope = self.kv_a_layernorm(k_nope)
|
k_nope = self.kv_a_layernorm(k_nope)
|
||||||
|
|||||||
Reference in New Issue
Block a user