#010

fused_qk_rmsnorm

fp16 rtp-llm · both · rtp_llm.cpp.kernels.rocm.fused_qk_rmsnorm · importance 1.9%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self, head_num: int, kv_head_num: int, size_per_head: int=128, eps: float=1e-06) -> None:
        super().__init__()
        self.head_num = head_num
        self.kv_head_num = kv_head_num
        self.size_per_head = size_per_head
        self.eps = eps
        self.q_size = head_num * size_per_head
        self.kv_size = kv_head_num * size_per_head

    def _rmsnorm(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
        input_dtype = x.dtype
        x = x.to(torch.float32)
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return weight * x.to(input_dtype)

    def forward(self, hidden_states: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor) -> torch.Tensor:
        (q, k, v) = hidden_states.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q = self._rmsnorm(q.reshape(-1, self.size_per_head), q_weight).view(q.shape)
        k = self._rmsnorm(k.reshape(-1, self.size_per_head), k_weight).view(k.shape)
        return torch.cat([q, k, v], dim=-1)

Shapes

T_SOL hardware:

#	num_tokens	head_num	kv_head_num	size_per_head	dtype	T_SOL(XPU-A)	T_Prod	S
0	1	64	8	128	fp16	0.01 us	7.70 us	0.1%
1	1	32	8	128	fp16	0.00 us	7.80 us	0.0%
3	1	32	4	128	fp16	0.00 us	7.80 us	0.0%
4	1	40	8	128	fp16	0.01 us	7.70 us	0.1%

Input Generation

input.py

import torch

def _make_inputs(num_tokens: int, head_num: int, kv_head_num: int, size_per_head: int=128, dtype: str='fp16') -> dict[str, torch.Tensor]:
    dt = torch.float16 if dtype == 'fp16' else torch.bfloat16
    hidden_size = head_num * size_per_head + 2 * kv_head_num * size_per_head
    hidden_states = torch.randn(num_tokens, hidden_size, dtype=dt, device='cuda')
    q_weight = torch.randn(size_per_head, dtype=dt, device='cuda')
    k_weight = torch.randn(size_per_head, dtype=dt, device='cuda')
    return {'hidden_states': hidden_states, 'q_weight': q_weight, 'k_weight': k_weight}