fix bugs with triton3.4.0 (#3946)

grimoire · web-flow · commit a7d7a581cbb1 · 2025-09-11T22:01:04.000+08:00
* fix bugs with triton3.4.0

* random seed

* bug fixing
diff --git a/lmdeploy/pytorch/check_env/triton.py b/lmdeploy/pytorch/check_env/triton.py
@@ -3,7 +3,7 @@
 
 from .base import BaseChecker
 
-MAX_TRITON_VERSION = '3.3.1'
+MAX_TRITON_VERSION = '3.4.0'
 MIN_TRITON_VERSION = '3.0.0'
 
 
diff --git a/lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py b/lmdeploy/pytorch/kernels/cuda/apply_rotary_pos_emb.py
@@ -44,7 +44,7 @@ def apply_rotary_pos_emb_qk_kernel(
     feat_mask = feat_offset_l < half_size
     feat_offset_l = feat_offset_l % half_size
     feat_offset_h = half_size + feat_offset_l
-    seq_mask = pos_mask[:, None] and feat_mask[None, :]
+    seq_mask = pos_mask[:, None] & feat_mask[None, :]
     cs_offset_l = pos_offset[:, None] * feat_size + feat_offset_l[None, :]
     cs_offset_h = pos_offset[:, None] * feat_size + feat_offset_h[None, :]
     q_elem_type = Q.dtype.element_ty
diff --git a/lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py b/lmdeploy/pytorch/kernels/cuda/blocked_fp8_fused_moe.py
@@ -125,7 +125,7 @@ def fused_moe_blocked_f8_kernel(
     k_start = BLOCK_SIZE_K
     offs_ksa = k_start // group_ak
     offs_ksb = k_start // group_bk
-    a_scale = tl.load(as_ptrs + offs_ksa * stride_ask, mask=mask_sid and k_start < K, other=1.0)
+    a_scale = tl.load(as_ptrs + offs_ksa * stride_ask, mask=mask_sid & (k_start < K), other=1.0)
     b_scale = tl.load(bs_ptrs + offs_ksb * stride_bsk, mask=k_start < K, other=1.0)
     acc_scale1 = tl.maximum(a_scale * b_scale, 1e-12)
     acc_ratio = acc_scale0 / acc_scale1
@@ -136,7 +136,7 @@ def fused_moe_blocked_f8_kernel(
         k_start = (k + 2) * BLOCK_SIZE_K
         offs_ksa = k_start // group_ak
         offs_ksb = k_start // group_bk
-        a_scale = tl.load(as_ptrs + offs_ksa * stride_ask, mask=mask_sid and k_start < K, other=1.0)
+        a_scale = tl.load(as_ptrs + offs_ksa * stride_ask, mask=mask_sid & (k_start < K), other=1.0)
         b_scale = tl.load(bs_ptrs + offs_ksb * stride_bsk, mask=k_start < K, other=1.0)
 
         # load ab
diff --git a/lmdeploy/pytorch/kernels/cuda/blocked_gemm_fp8.py b/lmdeploy/pytorch/kernels/cuda/blocked_gemm_fp8.py
@@ -132,15 +132,23 @@ def quant_fp8_tma(A: Tensor, group_size: int, dtype: torch.dtype = torch.float8_
     return _quant_fp8_launcher(A, group_size, out, scales)
 
 
+def _gemm_fp8_tma_pre_hook(nargs):
+    BLOCK_M = nargs['BLOCK_M']
+    BLOCK_N = nargs['BLOCK_N']
+    BLOCK_K = nargs['BLOCK_K']
+    nargs['desc_a'].block_shape = (BLOCK_M, BLOCK_K)
+    nargs['desc_b'].block_shape = (BLOCK_N, BLOCK_K)
+
+
 @triton.autotune(configs=[
     triton.Config({
         'BLOCK_M': 128,
         'BLOCK_N': 128,
-    }, num_stages=3, num_warps=8),
+    }, num_stages=3, num_warps=8, pre_hook=_gemm_fp8_tma_pre_hook),
     triton.Config({
         'BLOCK_M': 128,
         'BLOCK_N': 64,
-    }, num_stages=3, num_warps=4)
+    }, num_stages=3, num_warps=4, pre_hook=_gemm_fp8_tma_pre_hook)
 ],
                  key=['N', 'K'])
 @triton.jit
@@ -162,7 +170,6 @@ def _gemm_fp8_tma_kernel(
     stride_bsn: tl.constexpr,
     stride_cm,
     stride_cn: tl.constexpr,
-    dtype: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_K: tl.constexpr,
@@ -200,8 +207,8 @@ def _gemm_fp8_tma_kernel(
         b_scale = tl.load(bs_ptrs + offs_ksb * stride_bsk, mask=k_start < K, other=1.0)
 
         # load ab
-        a = tl._experimental_descriptor_load(desc_a, [off_m, off_k], [BLOCK_M, BLOCK_K], dtype)
-        b = tl._experimental_descriptor_load(desc_b, [off_n, off_k], [BLOCK_N, BLOCK_K], dtype).T
+        a = desc_a.load([off_m, off_k])
+        b = desc_b.load([off_n, off_k]).T
 
         # mma
         accumulator = tl.dot(a, b, acc=accumulator * acc_ratio[:, None])
@@ -348,42 +355,18 @@ def grid(META):
 
     # run_tma = False
     if run_tma:
-        from .utils import TmaAutoTuneHelper
+        from .utils import TensorDescriptor
 
-        desc_helper = TmaAutoTuneHelper()
-        desc_helper.init_tma_descriptor('desc_a')
-        desc_helper.init_tma_descriptor('desc_b')
-
-        desc_a = desc_helper.get_tma_descriptor_kernel_param('desc_a')
-        desc_b = desc_helper.get_tma_descriptor_kernel_param('desc_b')
+        dummy_block = (1, 1)
+        desc_a = TensorDescriptor.from_tensor(A, block_shape=dummy_block)
+        desc_b = TensorDescriptor.from_tensor(B.T, block_shape=dummy_block)
 
         def _grid_tma(META):
             """Grid tma."""
             BLOCK_M = META['BLOCK_M']
             BLOCK_N = META['BLOCK_N']
-            desc_helper.fill_2d_tma_descriptor('desc_a',
-                                               A.data_ptr(),
-                                               dim1=M,
-                                               dim0=K,
-                                               block_dim1=BLOCK_M,
-                                               block_dim0=BLOCK_K,
-                                               element_size=A.element_size())
-            desc_helper.fill_2d_tma_descriptor('desc_b',
-                                               B.data_ptr(),
-                                               dim1=N,
-                                               dim0=K,
-                                               block_dim1=BLOCK_N,
-                                               block_dim0=BLOCK_K,
-                                               element_size=B.element_size())
             return (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), )
 
-        if A.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz):
-            dtype = tl.float8e4nv
-        elif A.dtype in (torch.float8_e5m2, torch.float8_e5m2fnuz):
-            dtype = tl.float8e5
-        else:
-            raise RuntimeError(f'Not supported dtype: {A.dtype}')
-
         _gemm_fp8_tma_kernel[_grid_tma](
             desc_a,
             A_scale,
@@ -402,13 +385,8 @@ def _grid_tma(META):
             stride_bsn=B_scale.stride(1),
             stride_cm=C.stride(0),
             stride_cn=C.stride(1),
-            dtype=dtype,
-            # BLOCK_M=BLOCK_M,
-            # BLOCK_N=BLOCK_N,
             BLOCK_K=BLOCK_K,
             GROUP_M=8,
-            # num_warps=num_warps,
-            # num_stages=num_stages,
         )
     else:
         _gemm_fp8_kernel[grid](
diff --git a/lmdeploy/pytorch/kernels/cuda/flashattention.py b/lmdeploy/pytorch/kernels/cuda/flashattention.py
@@ -79,7 +79,7 @@ def _prefill_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, q1, k1_ptrs, loop_start
             qk = qk * tl_log2(math.e)
             qk_mask = (history_mask[:, None]) >= (start_n + offs_n[None, :])
             if window_size > 0:
-                qk_mask = qk_mask and ((start_n + offs_n[None, :]) >= kv_min_loc[:, None])
+                qk_mask = qk_mask & ((start_n + offs_n[None, :]) >= kv_min_loc[:, None])
             qk = tl.where(
                 qk_mask,
                 qk,
@@ -218,7 +218,7 @@ def _flash_prefill_fwd_kernel(
     offs_dk = tl.multiple_of(tl.max_contiguous(offs_dk % head_dim_k, BLOCK_DK), BLOCK_DK)
     off_q = ((q_start_loc + offs_m[:, None]) * stride_qs + head_id * stride_qh + offs_dk[None, :] * stride_qd)
     q_ptrs = q_ptr + off_q
-    q = tl.load(q_ptrs, mask=(offs_m[:, None] < q_seqlen and mask_dk[None, :]))
+    q = tl.load(q_ptrs, mask=((offs_m[:, None] < q_seqlen) & mask_dk[None, :]))
 
     k_ptrs = tl.make_block_ptr(
         base=k_ptr + kv_start_loc * stride_ks + kv_head_id * stride_kh,
@@ -252,7 +252,7 @@ def _flash_prefill_fwd_kernel(
         offs_dk1 = tl.multiple_of(tl.max_contiguous(offs_dk1 % head_dim_k, BLOCK_DK1), BLOCK_DK1)
         offs_q1 = ((q_start_loc + offs_m[:, None]) * stride_qs + head_id * stride_qh + offs_dk1[None, :] * stride_qd)
         q1_ptrs = q_ptr + offs_q1
-        q1 = tl.load(q1_ptrs, mask=(offs_m[:, None] < q_seqlen and mask_dk1[None, :]))
+        q1 = tl.load(q1_ptrs, mask=((offs_m[:, None] < q_seqlen) & mask_dk1[None, :]))
         k1_ptrs = tl.make_block_ptr(
             base=k_ptr + kv_start_loc * stride_ks + kv_head_id * stride_kh,
             shape=(head_dim_k, kv_seqlen),
diff --git a/lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py b/lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py
@@ -74,10 +74,10 @@ def _flatten_kv_cache(
                offs_dv[None, :] * stride_vod)
 
     kc = tl.load(kc_ptrs)
-    tl.store(ko_ptrs, kc, mask=mask_bs[:, None] and mask_dk[None, :])
+    tl.store(ko_ptrs, kc, mask=mask_bs[:, None] & mask_dk[None, :])
     if HEAD_DIM_V > 0:
         vc = tl.load(vc_ptrs)
-        tl.store(vo_ptrs, vc, mask=mask_bs[:, None] and mask_dv[None, :])
+        tl.store(vo_ptrs, vc, mask=mask_bs[:, None] & mask_dv[None, :])
 
 
 @triton.jit
@@ -181,15 +181,15 @@ def _flatten_kv_cache_quant(
     kz = tl.load(ksz_ptrs + stride_kszd)
     ksz = ks * kz
     kq = (kc * ks[:, None] - ksz[:, None]).to(ko_ptr.dtype.element_ty)
-    tl.store(ko_ptrs, kq, mask=mask_bs[:, None] and mask_dok[None, :])
+    tl.store(ko_ptrs, kq, mask=mask_bs[:, None] & mask_dok[None, :])
     vc = tl.load(vc_ptrs)
     if quant_policy == 4:
         vc = _dequant_int4(vc, HEAD_DIM_V, BLOCK_DV)
     vs = tl.load(vsz_ptrs)
     vz = tl.load(vsz_ptrs + stride_vszd)
     vsz = vs * vz
     vq = (vc * vs[:, None] - vsz[:, None]).to(vo_ptr.dtype.element_ty)
-    tl.store(vo_ptrs, vq, mask=mask_bs[:, None] and mask_dov[None, :])
+    tl.store(vo_ptrs, vq, mask=mask_bs[:, None] & mask_dov[None, :])
 
 
 def flatten_kv_cache(k_caches: Tensor,
diff --git a/lmdeploy/pytorch/kernels/cuda/pagedattention.py b/lmdeploy/pytorch/kernels/cuda/pagedattention.py
@@ -170,7 +170,7 @@ def _fwd_grouped_split_kernel(
         if start_n + BLOCK_N > history_len or window_size > 0:
             qk_mask = history_len >= (start_n + offs_n)
             if window_size > 0:
-                qk_mask = qk_mask and ((start_n + offs_n) >= kv_min_loc)
+                qk_mask = qk_mask & ((start_n + offs_n) >= kv_min_loc)
             qk = tl.where(
                 qk_mask[None, :],
                 qk,
@@ -388,7 +388,7 @@ def _fwd_grouped_split_quant_kernel(
         if start_n + BLOCK_N > history_len or window_size > 0:
             qk_mask = history_len >= (start_n + offs_n)
             if window_size > 0:
-                qk_mask = qk_mask and ((start_n + offs_n) >= kv_min_loc)
+                qk_mask = qk_mask & ((start_n + offs_n) >= kv_min_loc)
             qk = tl.where(
                 qk_mask[None, :],
                 qk,
diff --git a/lmdeploy/pytorch/kernels/cuda/utils.py b/lmdeploy/pytorch/kernels/cuda/utils.py
@@ -29,6 +29,8 @@
     (12, 0): 24,
 }
 
+TRITON_VERSION = version.parse(triton.__version__)
+
 
 @functools.lru_cache
 def get_device_props(device=None):
@@ -57,86 +59,9 @@ def supports_tma():
     if not ret:
         return False
 
-    TRITON_VERSION = version.parse(triton.__version__)
-    VALID_VERSION = version.parse('3.2.0')
-    return TRITON_VERSION >= VALID_VERSION
-
-
-# Copy from:
-# https://github.com/triton-lang/triton/blob/main/python/triton/tools/experimental_descriptor.py
-class TmaDescKernelParam:
-    TMA_DESC_SIZE = 128
-
-    def __init__(self):
-        self.desc = torch.empty(self.TMA_DESC_SIZE, dtype=torch.uint8, device='cpu')
-
-    def fill_(self, ptr, dims, block_dims, element_size):
-        assert len(dims) == len(block_dims)
-        assert 1 <= len(dims) <= 2
-        assert self.desc.data_ptr() % 64 == 0
-
-        if len(dims) == 1:
-            triton.runtime.driver.active.utils.fill_1d_tma_descriptor(ptr, dims[0], block_dims[0], element_size,
-                                                                      self.desc.data_ptr())
-        else:
-            triton.runtime.driver.active.utils.fill_2d_tma_descriptor(ptr, dims[0], dims[1], block_dims[0],
-                                                                      block_dims[1], element_size, self.desc.data_ptr())
-
-    # Return a CUtensorMap* pointer in host memory
-    def tma_desc_cpu_ptr(self):
-        return self.desc.data_ptr()
-
-
-# Copy from:
-# https://github.com/triton-lang/triton/blob/main/python/triton/tools/experimental_descriptor.py
-def create_1d_tma_descriptor_custom(ptr, dim, block_dim, element_size):
-    desc = TmaDescKernelParam()
-    desc.fill_(ptr, [dim], [block_dim], element_size)
-    return desc
-
-
-# Copy from:
-# https://github.com/triton-lang/triton/blob/main/python/triton/tools/experimental_descriptor.py
-def create_2d_tma_descriptor_custom(ptr, dim1, dim0, block_dim1, block_dim0, element_size):
-    desc = TmaDescKernelParam()
-    desc.fill_(ptr, [dim1, dim0], [block_dim1, block_dim0], element_size)
-    return desc
-
-
-try:
-    from triton.tools.experimental_descriptor import create_1d_tma_descriptor, create_2d_tma_descriptor  # noqa
-except BaseException:
-    create_1d_tma_descriptor = create_1d_tma_descriptor_custom
-    create_2d_tma_descriptor = create_2d_tma_descriptor_custom
-
-
-class TmaAutoTuneHelper:
-
-    # duck typing wrapper to implement the same interface as TmaDescKernelParam in Triton PR #4498
-    class KernelParamWrapper:
-
-        def __init__(self, desc):
-            self.desc = desc
-
-        def tma_desc_cpu_ptr(self):
-            return self.desc.data_ptr()
-
-    TMA_SIZE = 128
-
-    def __init__(self):
-        self.fill_2d_tma_descriptor_inner = (triton.runtime.driver.active.utils.fill_2d_tma_descriptor)
-        self.descriptors = {}
-
-    # Call this method outside of the lambda function for grid size
-    def init_tma_descriptor(self, name):
-        self.descriptors[name] = torch.empty(TmaAutoTuneHelper.TMA_SIZE, device='cpu', dtype=torch.int8)
+    VALID_VERSION = version.parse('3.4.0')
+    return TRITON_VERSION == VALID_VERSION
 
-    # Call this method inside the lambda function for grid size
-    def fill_2d_tma_descriptor(self, name, ptr, dim1, dim0, block_dim1, block_dim0, element_size):
-        desc_x = self.descriptors[name]
-        assert desc_x.data_ptr() % 64 == 0
-        self.fill_2d_tma_descriptor_inner(ptr, dim1, dim0, block_dim1, block_dim0, element_size, desc_x.data_ptr())
 
-    def get_tma_descriptor_kernel_param(self, name):
-        assert self.descriptors[name] is not None
-        return self.KernelParamWrapper(self.descriptors[name])
+if supports_tma():
+    from triton.tools.tensor_descriptor import TensorDescriptor  # noqa: F401
diff --git a/lmdeploy/pytorch/messages.py b/lmdeploy/pytorch/messages.py
@@ -114,13 +114,18 @@ def from_gen_config(self, gen_config: GenerationConfig):
         logprobs = gen_config.logprobs
         if logprobs is None:
             logprobs = -1
+
+        random_seed = gen_config.random_seed
+        if random_seed is None:
+            import random
+            random_seed = random.getrandbits(64)
         return SamplingParam(top_p=top_p,
                              top_k=top_k,
                              min_p=min_p,
                              temperature=temperature,
                              repetition_penalty=repetition_penalty,
                              ignore_eos=gen_config.ignore_eos,
-                             random_seed=gen_config.random_seed,
+                             random_seed=random_seed,
                              stop_words=stop_words,
                              bad_words=bad_words,
                              response_format=response_format,
diff --git a/tests/pytorch/kernel/test_apply_rotary.py b/tests/pytorch/kernel/test_apply_rotary.py
@@ -94,5 +94,8 @@ def test_apply_rotary(self, q_states, k_states, cos, sin, gt):
         if q_states.dtype == torch.float16:
             rtol = 1e-5
             atol = 1e-3
+        elif q_states.dtype == torch.bfloat16:
+            rtol = 1e-5
+            atol = 1e-2
         torch.testing.assert_close(q_embed, q_gt, rtol=rtol, atol=atol)
         torch.testing.assert_close(k_embed, k_gt, rtol=rtol, atol=atol)
