Add ApplyRoPE and RMSNorm kernels written in OpenAI Triton.

config.json

@@ -44,6 +44,7 @@
   "use_cache": true,
   "use_dynamic_ntk": true,
   "use_flash_attn": "auto",
+  "use_triton": "auto",
   "use_logn_attn": true,
   "vocab_size": 151936
 }

configuration_qwen.py

@@ -32,6 +32,7 @@ class QWenConfig(PretrainedConfig):
         use_dynamic_ntk=True,
         use_logn_attn=True,
         use_flash_attn="auto",
+        use_triton="auto",
         intermediate_size=22016,
         no_bias=True,
         tie_word_embeddings=False,
@@ -61,6 +62,7 @@ class QWenConfig(PretrainedConfig):
         self.use_dynamic_ntk = use_dynamic_ntk
         self.use_logn_attn = use_logn_attn
         self.use_flash_attn = use_flash_attn
+        self.use_triton = use_triton
         self.no_bias = no_bias
         self.use_cache_quantization = use_cache_quantization
         self.use_cache_kernel = use_cache_kernel

modeling_qwen.py

@@ -35,7 +35,7 @@ except ImportError:
 from torch import nn
 
 SUPPORT_CUDA = torch.cuda.is_available()
-SUPPORT_BF16 = SUPPORT_CUDA and torch.cuda.is_bf16_supported()
+SUPPORT_BF16 = SUPPORT_CUDA and torch.cuda.get_device_capability(0)[0] >= 8
 SUPPORT_FP16 = SUPPORT_CUDA and torch.cuda.get_device_capability(0)[0] >= 7
 SUPPORT_TORCH2 = hasattr(torch, '__version__') and int(torch.__version__.split(".")[0]) >= 2
 
@@ -76,7 +76,9 @@ We detect you have activated flash attention support, but running model computat
 """
 
 apply_rotary_emb_func = None
+apply_rotary_emb_func_triton = None
 rms_norm = None
+rms_norm_triton = None
 flash_attn_unpadded_func = None
 flash_attn_func = None
 
@@ -120,6 +122,24 @@ def _import_flash_attn():
             "https://github.com/Dao-AILab/flash-attention"
         )
 
+def _import_triton():
+    global apply_rotary_emb_func_triton, rms_norm_triton
+    try:
+        from .triton_kernels import apply_rotary_emb as __apply_rotary_emb, rms_norm as __rms_norm
+        if apply_rotary_emb_func is not None:
+            logger.warn(
+                "Using Triton rotary kernel instead of flash_attn for inference."
+            )
+        apply_rotary_emb_func_triton = __apply_rotary_emb
+        if rms_norm is not None:
+            logger.warn(
+                "Using Triton rms_norm kernel instead of flash_attn for inference."
+            )
+        rms_norm_triton = __rms_norm
+    except ImportError:
+        logger.warn("Warning: Failed to import Triton kernels.")
+        return
+
 def quantize_cache_v(fdata, bits, qmax, qmin):
     # b, s, head, h-dim->b, head, s, h-dim
     qtype = torch.uint8
@@ -978,6 +998,12 @@ class QWenLMHeadModel(QWenPreTrainedModel):
         if config.use_flash_attn:
             _import_flash_attn()
 
+        if config.use_triton == "auto":
+            logger.warn("Try importing Triton kernels for faster inference...")
+            config.use_triton = SUPPORT_TORCH2
+        if config.use_triton:
+            _import_triton()
+
         self.transformer = QWenModel(config)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
 
@@ -1335,7 +1361,9 @@ def apply_rotary_pos_emb(t, freqs):
     rot_dim = freqs[0].shape[-1]
     cos, sin = freqs
     t_float = t.float()
-    if apply_rotary_emb_func is not None and t.is_cuda:
+    if apply_rotary_emb_func_triton is not None and t.is_cuda and (not t.requires_grad):
+        return apply_rotary_emb_func_triton(t, cos, sin)
+    elif apply_rotary_emb_func is not None and t.is_cuda:
         # apply_rotary_emb in flash_attn requires cos/sin to be of 
         # shape (seqlen, rotary_dim / 2) and apply rotary embedding 
         # to the first rotary_dim of the input
@@ -1358,7 +1386,9 @@ class RMSNorm(torch.nn.Module):
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
 
     def forward(self, x):
-        if rms_norm is not None and x.is_cuda:
+        if rms_norm_triton is not None and x.is_cuda and (not x.requires_grad):
+            return rms_norm_triton(x, self.weight, self.eps)
+        elif rms_norm is not None and x.is_cuda:
             return rms_norm(x, self.weight, self.eps)
         else:
             output = self._norm(x.float()).type_as(x)

triton_kernels.py

@@ -0,0 +1,115 @@
+from typing import Any, Callable, Dict, Hashable, Tuple
+
+import torch
+import triton
+import triton.language as tl
+from triton.compiler import CompiledKernel
+from triton.runtime import JITFunction
+
+try:
+    import triton.language.math as tlmath  # Triton 2.1
+except ImportError:
+    import triton.language.libdevice as tlmath  # Triton 2.0
+
+
+class TritonKernel:
+    def __init__(
+        self,
+        kernel_fn: JITFunction,
+        grid_fn: Callable[[Tuple[Any, ...]], Tuple[int, int, int]],
+    ) -> None:
+        self.kernel_fn_ = kernel_fn
+        self.grid_fn_ = grid_fn
+        self.kernel_cache_: Dict[Hashable, CompiledKernel] = {}
+
+    def run(self, *args, **kwargs):
+        # Set current device
+        input_device = args[0].device
+        prev_dev_idx, cur_dev_idx = -1, torch.cuda.current_device()
+        if input_device.index != cur_dev_idx:
+            prev_dev_idx = cur_dev_idx
+            torch.cuda.set_device(input_device.index)
+
+        # Compute grid
+        grid = self.grid_fn_(args)
+
+        # Use cached kernel if possible
+        kernel_key = (input_device,) + tuple(kwargs.items())
+        if kernel_key in self.kernel_cache_:
+            kernel = self.kernel_cache_[kernel_key]
+            kernel[grid](*args)
+        else:
+            # Compile and store new kernel
+            kernel = self.kernel_fn_[grid](*args, **kwargs)
+            self.kernel_cache_[kernel_key] = kernel
+
+        # Restore previous device
+        torch.cuda.set_device(prev_dev_idx)
+
+
+@triton.jit
+def _apply_rope_fwd_kernel(X, Cos, Sin, Y, HEAD_DIM: tl.constexpr):
+    batch_idx, tok_idx, head_idx = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    seq_len, num_heads = tl.num_programs(1), tl.num_programs(2)
+    block_idx = tl.arange(0, HEAD_DIM)
+    x_base_idx = ((batch_idx * seq_len + tok_idx) * num_heads * 3 + head_idx) * HEAD_DIM
+    x = tl.load(X + x_base_idx + block_idx)
+    freq_idx = tok_idx * HEAD_DIM + block_idx
+    cos = tl.load(Cos + freq_idx)
+    rot_idx = (HEAD_DIM // 2 + block_idx) % HEAD_DIM
+    x_rot = tl.load(X + x_base_idx + rot_idx)
+    x_rot = tl.where(block_idx >= HEAD_DIM // 2, x_rot, -x_rot)
+    sin = tl.load(Sin + freq_idx)
+    y_idx = (
+        (batch_idx * seq_len + tok_idx) * num_heads + head_idx
+    ) * HEAD_DIM + block_idx
+    y = x * cos + x_rot * sin
+    tl.store(Y + y_idx, y.to(Y.dtype.element_ty))
+
+
+apply_rope_fwd_kernel = TritonKernel(
+    _apply_rope_fwd_kernel, lambda args: tuple(args[0].shape[:3])
+)
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
+    y = torch.empty(x.shape, dtype=x.dtype, device=x.device)
+    apply_rope_fwd_kernel.run(x, cos, sin, y, HEAD_DIM=x.size(-1))
+    return y
+
+
+@triton.jit
+def _rms_norm_fwd_kernel(X, W, Y, eps, hidden_dim, BLOCK_SIZE: tl.constexpr):
+    tok_idx = tl.program_id(0)
+
+    mean_sq = tl.zeros([BLOCK_SIZE], tl.float32)
+    for offset in range(0, hidden_dim, BLOCK_SIZE):
+        dim_idx = offset + tl.arange(0, BLOCK_SIZE)
+        x = tl.load(
+            X + tok_idx * hidden_dim + dim_idx, mask=dim_idx < hidden_dim, other=0
+        ).to(tl.float32)
+        mean_sq += x * x / hidden_dim
+    rrms = tlmath.rsqrt(tl.sum(mean_sq, 0) + eps)
+
+    for offset in range(0, hidden_dim, BLOCK_SIZE):
+        dim_idx = offset + tl.arange(0, BLOCK_SIZE)
+        dim_mask = dim_idx < hidden_dim
+        hidden_idx = tok_idx * hidden_dim + dim_idx
+        x = tl.load(X + hidden_idx, mask=dim_mask, other=0)
+        w = tl.load(W + dim_idx, mask=dim_mask, other=0)
+        y = x * rrms * w
+        tl.store(Y + hidden_idx, y.to(Y.dtype.element_ty), mask=dim_mask)
+
+
+rms_norm_fwd_kernel = TritonKernel(
+    _rms_norm_fwd_kernel, lambda args: (args[0].shape[:-1].numel(), 1, 1)
+)
+
+
+def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float):
+    y = torch.empty_like(x)
+    hidden_dim = x.size(-1)
+    rms_norm_fwd_kernel.run(
+        x, weight, y, eps, hidden_dim, BLOCK_SIZE=triton.next_power_of_2(hidden_dim)
+    )
+    return y