Improve performance witih Triton 2.0 and adapt to latest Qwen releases.

modeling_qwen.py

@@ -78,6 +78,7 @@ 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
 
@@ -122,28 +123,22 @@ def _import_flash_attn():
         )
 
 def _import_triton():
-    global apply_rotary_emb_func_triton, rms_norm
+    global apply_rotary_emb_func_triton, rms_norm_triton
     try:
-        from .triton_kernels import triton, apply_rotary_emb as __apply_rotary_emb, rms_norm as __rms_norm
+        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(
-                "rotary kernel imported from flash_attn is replaced by Triton kernel."
+                "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(
-                "rms_norm kernel imported from flash_attn is replaced by Triton kernel."
+                "Using Triton rms_norm kernel instead of flash_attn for inference."
             )
-        rms_norm = __rms_norm
+        rms_norm_triton = __rms_norm
     except ImportError:
         logger.warn("Warning: Failed to import Triton kernels.")
         return
-    
-    if int(triton.__version__.split(".")[1]) == 0:
-        logger.warn(
-            "Triton 2.0 is detected in your environment. It is recommended that you upgrade to Triton 2.1 by "
-            "`pip install triton==2.1` for better performance if you do not use TorchInductor in PyTorch 2.0."
-        )
 
 def quantize_cache_v(fdata, bits, qmax, qmin):
     # b, s, head, h-dim->b, head, s, h-dim
@@ -1004,9 +999,9 @@ class QWenLMHeadModel(QWenPreTrainedModel):
             _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:
-            logger.warn("Try importing Triton kernels for faster inference...")
             _import_triton()
 
         self.transformer = QWenModel(config)
@@ -1366,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
@@ -1389,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

@@ -1,13 +1,10 @@
-from functools import partial
 from typing import Any, Callable, Dict, Hashable, Tuple
 
 import torch
 import triton
 import triton.language as tl
-from torch.autograd.function import Function, FunctionCtx
 from triton.compiler import CompiledKernel
-from triton.runtime import KernelInterface
-from triton.runtime.autotuner import Autotuner
+from triton.runtime import JITFunction
 
 try:
     import triton.language.math as tlmath  # Triton 2.1
@@ -18,10 +15,10 @@ except ImportError:
 class TritonKernel:
     def __init__(
         self,
-        kernel_fn_: KernelInterface,
+        kernel_fn: JITFunction,
         grid_fn: Callable[[Tuple[Any, ...]], Tuple[int, int, int]],
     ) -> None:
-        self.kernel_fn_ = kernel_fn_
+        self.kernel_fn_ = kernel_fn
         self.grid_fn_ = grid_fn
         self.kernel_cache_: Dict[Hashable, CompiledKernel] = {}
 
@@ -37,24 +34,14 @@ class TritonKernel:
         grid = self.grid_fn_(args)
 
         # Use cached kernel if possible
-        kernel_key = (input_device,)
-        if isinstance(self.kernel_fn_, Autotuner):
-            kernel_key += tuple(args[ki] for ki in self.kernel_fn_.key_idx)
-        else:
-            kernel_key += tuple(kwargs.items())
+        kernel_key = (input_device,) + tuple(kwargs.items())
         if kernel_key in self.kernel_cache_:
             kernel = self.kernel_cache_[kernel_key]
             kernel[grid](*args)
-            return
-
-        # Compile new kernel
-        if isinstance(self.kernel_fn_, Autotuner):
-            kernel = self.kernel_fn_[grid](*args)
         else:
+            # Compile and store new kernel
             kernel = self.kernel_fn_[grid](*args, **kwargs)
-
-        # Store kernel
-        self.kernel_cache_[kernel_key] = kernel
+            self.kernel_cache_[kernel_key] = kernel
 
         # Restore previous device
         torch.cuda.set_device(prev_dev_idx)
@@ -77,7 +64,7 @@ def _apply_rope_fwd_kernel(X, Cos, Sin, Y, HEAD_DIM: tl.constexpr):
         (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)
+    tl.store(Y + y_idx, y.to(Y.dtype.element_ty))
 
 
 apply_rope_fwd_kernel = TritonKernel(
@@ -85,28 +72,12 @@ apply_rope_fwd_kernel = TritonKernel(
 )
 
 
-class ApplyRotaryEmb(Function):
-    @staticmethod
-    def forward(
-        ctx: FunctionCtx, 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
-
-
 def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
-    return ApplyRotaryEmb.apply(x, cos, sin)
+    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.autotune(
-    configs=[
-        triton.Config({"BLOCK_SIZE": 4096}),
-        triton.Config({"BLOCK_SIZE": 2048}),
-        triton.Config({"BLOCK_SIZE": 1024}),
-    ],
-    key=[],
-)
 @triton.jit
 def _rms_norm_fwd_kernel(X, W, Y, eps, hidden_dim, BLOCK_SIZE: tl.constexpr):
     tok_idx = tl.program_id(0)
@@ -135,13 +106,10 @@ rms_norm_fwd_kernel = TritonKernel(
 )
 
 
-class RMSNorm(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx: FunctionCtx, x: torch.Tensor, w: torch.Tensor, eps: float):
-        y = torch.empty_like(x)
-        rms_norm_fwd_kernel.run(x, w, y, eps, x.size(-1))
-        return y
-
-
 def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float):
-    return RMSNorm.apply(x, weight, eps)
+    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