[sgl-kernel][1/2] Fused qk_norm_rope for Qwen3-MoE (#14036)
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
@@ -251,6 +251,42 @@ def apply_shuffle_mul_sum(
|
||||
)
|
||||
|
||||
|
||||
def fused_qk_norm_rope(
|
||||
qkv: torch.Tensor,
|
||||
num_heads_q: int,
|
||||
num_heads_k: int,
|
||||
num_heads_v: int,
|
||||
head_dim: int,
|
||||
eps: float,
|
||||
q_weight: torch.Tensor,
|
||||
k_weight: torch.Tensor,
|
||||
base: float,
|
||||
is_neox: bool,
|
||||
position_ids: torch.Tensor,
|
||||
factor: float,
|
||||
low: float,
|
||||
high: float,
|
||||
attention_factor: float,
|
||||
) -> None:
|
||||
torch.ops.sgl_kernel.fused_qk_norm_rope(
|
||||
qkv,
|
||||
num_heads_q,
|
||||
num_heads_k,
|
||||
num_heads_v,
|
||||
head_dim,
|
||||
eps,
|
||||
q_weight,
|
||||
k_weight,
|
||||
base,
|
||||
is_neox,
|
||||
position_ids,
|
||||
factor,
|
||||
low,
|
||||
high,
|
||||
attention_factor,
|
||||
)
|
||||
|
||||
|
||||
def cutlass_fp4_group_mm(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
|
||||
Reference in New Issue
Block a user