liger-kernel-nightly 0.6.4.dev20251229100549__tar.gz → 0.6.4.dev20260113013032__tar.gz

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/amd-ci.yml

@@ -26,10 +26,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 
@@ -50,10 +50,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/benchmark.yml

@@ -35,22 +35,26 @@ jobs:
       OUTPUT_DIR: benchmarks
       OUTPUT_FILENAME: benchmark.csv
       GENERATED_CSV: benchmark/data/all_benchmark_data.csv
+      # Sanitize user-controlled inputs by declaring them as environment variables
+      # This prevents command injection attacks by filtering dangerous characters
+      INPUT_COMMIT_HASH: ${{ github.event.inputs.commit_hash }}
+      INPUT_OVERWRITE: ${{ github.event.inputs.overwrite }}
 
 
     steps:
       # Step: Decide the commit hash to use
       # Step: Checkout full history so we can check out any commit
       - name: Checkout full repo history
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           fetch-depth: 0  # Important: so we can checkout arbitrary commit
           
       - name: Determine commit hash to checkout
         id: choose_commit
         run: |
-          if [ "${{ github.event_name}}" == "workflow_dispatch" ] && [ "${{ github.event.inputs.commit_hash }}" != "main" ]; then
-            echo "Using manual input commit: ${{ github.event.inputs.commit_hash }}"
-            echo "hash=${{ github.event.inputs.commit_hash }}" >> $GITHUB_OUTPUT
+          if [ "${{ github.event_name}}" == "workflow_dispatch" ] && [ "$INPUT_COMMIT_HASH" != "main" ]; then
+            echo "Using manual input commit: $INPUT_COMMIT_HASH"
+            echo "hash=$INPUT_COMMIT_HASH" >> $GITHUB_OUTPUT
           else
             echo "Using latest commit from main"
             echo "hash=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT
@@ -60,10 +64,10 @@ jobs:
       - name: Replace benchmark folder from main (manual only, commit ≠ main)
         if: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.commit_hash != 'main' }}
         run: |
-          echo "Detected manual trigger with commit_hash = ${{ github.event.inputs.commit_hash }}"
+          echo "Detected manual trigger with commit_hash = $INPUT_COMMIT_HASH"
           
           # Save current branch (detached HEAD at old commit)
-          ORIG_COMMIT=${{ github.event.inputs.commit_hash }}
+          ORIG_COMMIT="$INPUT_COMMIT_HASH"
 
           # Fetch and checkout main
           git fetch origin main
@@ -72,7 +76,7 @@ jobs:
           # Save benchmark folder from main
           cp -r benchmark /tmp/benchmark_main
           # Checkout back to target commit
-          git checkout $ORIG_COMMIT
+          git checkout "$ORIG_COMMIT"
           # Replace old benchmark with one from main
           rm -rf benchmark
           cp -r /tmp/benchmark_main benchmark
@@ -85,7 +89,7 @@ jobs:
           
           if curl --output /dev/null --silent --head --fail "$BENCHMARK_URL"; then
             echo "Benchmark already exists for commit $COMMIT_HASH"
-            if [ "${{ github.event.inputs.overwrite }}" != "true" ]; then
+            if [ "$INPUT_OVERWRITE" != "true" ]; then
               echo "Overwrite is false - exiting"
               exit 1
             else
@@ -96,7 +100,7 @@ jobs:
           fi
 
       - name: Set up Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v6
         with:
           python-version: '3.10'
 
@@ -117,7 +121,7 @@ jobs:
         
        # Step 5: Checkout gh-pages branch in a subfolderAdd commentMore actions
       - name: Checkout gh-pages
-        uses: actions/checkout@v3
+        uses: actions/checkout@v6
         with:
           ref: gh-pages
           path: gh-pages

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/docs.yml

@@ -13,16 +13,16 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v6
       - name: Configure Git Credentials
         run: |
           git config user.name github-actions[bot]
           git config user.email 41898282+github-actions[bot]@users.noreply.github.com
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@v6
         with:
           python-version: 3.x
       - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV 
-      - uses: actions/cache@v4
+      - uses: actions/cache@v5
         with:
           key: mkdocs-material-${{ env.cache_id }}
           path: .cache

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/intel-ci.yml

@@ -26,10 +26,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 
@@ -58,7 +58,7 @@ jobs:
         apt-get clean && rm -rf /var/lib/apt/lists/*
 
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
       
     - name: Setup Dependencies
       shell: bash

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/nvi-ci.yml

@@ -25,10 +25,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 
@@ -49,10 +49,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 
@@ -75,10 +75,10 @@ jobs:
 
     steps:
     - name: Checkout code
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/publish-nightly.yml

@@ -13,10 +13,10 @@ jobs:
 
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.8'
 

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/.github/workflows/publish-release.yml

@@ -10,10 +10,10 @@ jobs:
 
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v3
+      uses: actions/checkout@v6
 
     - name: Set up Python
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@v6
       with:
         python-version: '3.10'
 

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: liger_kernel_nightly
-Version: 0.6.4.dev20251229100549
+Version: 0.6.4.dev20260113013032
 Summary: Efficient Triton kernels for LLM Training
 License: BSD 2-CLAUSE LICENSE
         Copyright 2024 LinkedIn Corporation 

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "liger_kernel_nightly"
-version = "0.6.4.dev20251229100549"
+version = "0.6.4.dev20260113013032"
 description = "Efficient Triton kernels for LLM Training"
 urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
 readme = { file = "README.md", content-type = "text/markdown" }

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/src/liger_kernel/ops/backends/_ascend/ops/__init__.py

@@ -17,15 +17,33 @@ If __all__ is not defined, all public symbols will be auto-discovered.
 from liger_kernel.ops.backends._ascend.ops.geglu import LigerGELUMulFunction
 from liger_kernel.ops.backends._ascend.ops.geglu import geglu_backward
 from liger_kernel.ops.backends._ascend.ops.geglu import geglu_forward
+from liger_kernel.ops.backends._ascend.ops.qwen2vl_mrope import LigerQwen2VLMRopeFunction
+from liger_kernel.ops.backends._ascend.ops.qwen2vl_mrope import qwen2vl_mrope_backward
+from liger_kernel.ops.backends._ascend.ops.qwen2vl_mrope import qwen2vl_mrope_forward
 from liger_kernel.ops.backends._ascend.ops.rope import LigerRopeFunction
 from liger_kernel.ops.backends._ascend.ops.rope import rope_backward
 from liger_kernel.ops.backends._ascend.ops.rope import rope_forward
+from liger_kernel.ops.backends._ascend.ops.swiglu import LigerSiLUMulFunction
+from liger_kernel.ops.backends._ascend.ops.swiglu import swiglu_backward
+from liger_kernel.ops.backends._ascend.ops.swiglu import swiglu_forward
+from liger_kernel.ops.backends._ascend.ops.tvd import LigerTVDLossFunction
+from liger_kernel.ops.backends._ascend.ops.tvd import tv_distance_forward_triton
+from liger_kernel.ops.backends._ascend.ops.tvd import tvd_backward_triton
 
 __all__ = [
     "LigerGELUMulFunction",
     "geglu_forward",
     "geglu_backward",
+    "LigerQwen2VLMRopeFunction",
+    "qwen2vl_mrope_forward",
+    "qwen2vl_mrope_backward",
     "LigerRopeFunction",
     "rope_forward",
     "rope_backward",
+    "LigerSiLUMulFunction",
+    "swiglu_forward",
+    "swiglu_backward",
+    "LigerTVDLossFunction",
+    "tv_distance_forward_triton",
+    "tvd_backward_triton",
 ]

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/src/liger_kernel/ops/backends/_ascend/ops/geglu.py

@@ -130,20 +130,26 @@ def geglu_forward(a, b):
     dtype_size = a.element_size()
     # GEGLU forward tiling strategy:
     # - Calculates maximum safe block size based on UB capacity
-    # - Memory analysis:
-    #   * Inputs: a, b
-    #   * Intermediates: a_cubed, tanh_arg, tanh_result, geglu_a
-    #   * Output: c
-    #   * Total: ~7x * BLOCK_SIZE * dtype_size
-    # - Uses memory_multiplier=7.0 * BLOCK_SIZE * dtype_size * 8 bits for safety
+    # - Memory analysis (only buffers that occupy UB, excluding temporary variables):
+    #   * Inputs: a_row (4 bytes, float32), b_row (dtype_size bytes)
+    #   * Output: c_row (dtype_size bytes)
+    #   * Temporary variables (a_cubed, tanh_arg, tanh_result, geglu_a) are optimized to registers
+    #     and don't occupy UB since they are only used once
+    #   * For float16: a_row(4) + b_row(2) + c_row(2) = 8 bytes/element, ratio = 8/2 = 4.0
+    #   * For float32: a_row(4) + b_row(4) + c_row(4) = 12 bytes/element, ratio = 12/4 = 3.0
+    # - Uses memory_multiplier=4.0 (float16) or 3.0 (float32) * BLOCK_SIZE * dtype_size * 8 bits
     # - shapes: ((n_cols,),)
     # - tiling_dims: (0,) means first dimension can be tiled
     # - Returns: ((block_size,),)
     shapes = ((n_cols,),)
+    if dtype_size == 2:
+        memory_multiplier = 4.0
+    else:
+        memory_multiplier = 3.0
     tile_shapes = compute_default_tiling_strategy(
         safety_margin=0.80,
         dtype_size=dtype_size,
-        memory_multiplier=7.0,
+        memory_multiplier=memory_multiplier,
         shapes=shapes,
         tiling_dims=(0,),
     )
@@ -187,18 +193,34 @@ def geglu_backward(a, b, dc):
     dtype_size = dc.element_size()
     # GEGLU backward tiling strategy:
     # - Calculates maximum safe block size based on UB capacity
-    # - Memory analysis:
-    #   * More intermediates for gradient computation compared to forward
-    #   * Total: ~10x * BLOCK_SIZE * dtype_size
-    # - Uses memory_multiplier=10.0 * BLOCK_SIZE * dtype_size * 8 bits for safety
+    # - Memory analysis: Peak memory usage occurs when executing line 103 (term1 calculation)
+    #   At this point, the following buffers simultaneously occupy UB:
+    #   1. dc_row = tl.load(dc + col_offsets, ...)  # dtype_size bytes
+    #   2. a_row = tl.load(a + col_offsets, ...).to(tl.float32)  # 4 bytes (float32)
+    #   3. b_row = tl.load(b + col_offsets, ...)  # dtype_size bytes
+    #   4. tanh_result = tanh(tanh_arg)  # 4 bytes (float32), used in lines 95, 103, 104
+    #   5. geglu_a = 0.5 * a_row * (1 + tanh_result)  # 4 bytes (float32), used in lines 96, 98
+    #   6. db_row = dc_row.cast(tl.float32) * geglu_a  # 4 bytes (float32, computed at line 98, stored at line 109)
+    #   Note: term1 (line 103) is a temporary variable optimized to registers and doesn't occupy UB
+    #   Temporary variables (a_cubed, tanh_arg, term1, tanh_sq, term2) are optimized to registers
+    #   and don't occupy UB since they are only used once
+    #   * For float16: dc_row(2) + a_row(4) + b_row(2) + tanh_result(4) + geglu_a(4) + db_row(4)
+    #     = 20 bytes/element, ratio = 20/2 = 10.0
+    #   * For float32: dc_row(4) + a_row(4) + b_row(4) + tanh_result(4) + geglu_a(4) + db_row(4)
+    #     = 24 bytes/element, ratio = 24/4 = 6.0
+    # - Uses memory_multiplier=10.0 (float16) or 6.0 (float32) * BLOCK_SIZE * dtype_size * 8 bits
     # - shapes: ((n_cols,),)
     # - tiling_dims: (0,) means first dimension can be tiled
     # - Returns: ((block_size,),)
     shapes = ((n_cols,),)
+    if dtype_size == 2:
+        memory_multiplier = 10.0
+    else:
+        memory_multiplier = 6.0
     tile_shapes = compute_default_tiling_strategy(
         safety_margin=0.80,
         dtype_size=dtype_size,
-        memory_multiplier=10.0,
+        memory_multiplier=memory_multiplier,
         shapes=shapes,
         tiling_dims=(0,),
     )

liger_kernel_nightly-0.6.4.dev20260113013032/src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py

@@ -0,0 +1,285 @@
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.backends._ascend.ub_manager import compute_default_tiling_strategy
+
+
+@triton.jit
+def _triton_qwen2vl_mrope_npu(
+    q_ptr,
+    q_row_stride,
+    k_ptr,
+    k_row_stride,
+    cos,
+    sin,
+    sl,
+    bs: tl.constexpr,
+    n_qh: tl.constexpr,
+    n_kh: tl.constexpr,
+    hd: tl.constexpr,
+    mrope_section_t: tl.constexpr,
+    mrope_section_h: tl.constexpr,
+    BLOCK_Q: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BACKWARD_PASS: tl.constexpr = False,
+):
+    pid = tl.program_id(0).to(tl.int64)
+
+    t_end = mrope_section_t
+    h_end = t_end + mrope_section_h
+
+    t_cos = cos + pid * hd
+    h_cos = t_cos + bs * sl * hd
+    w_cos = h_cos + bs * sl * hd
+    t_sin = sin + pid * hd
+    h_sin = t_sin + bs * sl * hd
+    w_sin = h_sin + bs * sl * hd
+
+    q_base = q_ptr + pid * q_row_stride
+    k_base = k_ptr + pid * k_row_stride
+
+    d_idx = tl.arange(0, hd // 2)
+    d_mask = d_idx < (hd // 2)
+
+    pos_mask_t = d_idx < t_end
+    pos_mask_h = (d_idx >= t_end) & (d_idx < h_end)
+
+    text_cos_vals = tl.load(t_cos + d_idx, mask=d_mask, other=0)
+    text_sin_vals = tl.load(t_sin + d_idx, mask=d_mask, other=0)
+    height_cos_vals = tl.load(h_cos + d_idx, mask=d_mask, other=0)
+    height_sin_vals = tl.load(h_sin + d_idx, mask=d_mask, other=0)
+    width_cos_vals = tl.load(w_cos + d_idx, mask=d_mask, other=0)
+    width_sin_vals = tl.load(w_sin + d_idx, mask=d_mask, other=0)
+
+    cos_vals = tl.where(pos_mask_t, text_cos_vals, tl.where(pos_mask_h, height_cos_vals, width_cos_vals))
+    sin_vals = tl.where(pos_mask_t, text_sin_vals, tl.where(pos_mask_h, height_sin_vals, width_sin_vals))
+
+    for qh_block in range(0, n_qh, BLOCK_Q):
+        qh_idx = tl.arange(0, BLOCK_Q) + qh_block
+        qh_mask = qh_idx < n_qh
+
+        block_mask = qh_mask[:, None] & d_mask[None, :]
+        offsets = qh_idx[:, None] * hd + d_idx[None, :]
+
+        q_left = tl.load(q_base + offsets, mask=block_mask, other=0)
+        q_right = tl.load(q_base + offsets + (hd // 2), mask=block_mask, other=0)
+
+        if not BACKWARD_PASS:
+            new_left = q_left * cos_vals - q_right * sin_vals
+            new_right = q_right * cos_vals + q_left * sin_vals
+        else:
+            new_left = q_left * cos_vals + q_right * sin_vals
+            new_right = q_right * cos_vals - q_left * sin_vals
+
+        tl.store(q_base + offsets, new_left, mask=block_mask)
+        tl.store(q_base + offsets + (hd // 2), new_right, mask=block_mask)
+
+    for kh_block in range(0, n_kh, BLOCK_K):
+        kh_idx = tl.arange(0, BLOCK_K) + kh_block
+        kh_mask = kh_idx < n_kh
+
+        block_mask = kh_mask[:, None] & d_mask[None, :]
+        offsets = kh_idx[:, None] * hd + d_idx[None, :]
+
+        k_left = tl.load(k_base + offsets, mask=block_mask, other=0)
+        k_right = tl.load(k_base + offsets + (hd // 2), mask=block_mask, other=0)
+
+        if not BACKWARD_PASS:
+            new_left = k_left * cos_vals - k_right * sin_vals
+            new_right = k_right * cos_vals + k_left * sin_vals
+        else:
+            new_left = k_left * cos_vals + k_right * sin_vals
+            new_right = k_right * cos_vals - k_left * sin_vals
+
+        tl.store(k_base + offsets, new_left, mask=block_mask)
+        tl.store(k_base + offsets + (hd // 2), new_right, mask=block_mask)
+
+
+def qwen2vl_mrope_forward(q, k, cos, sin, mrope_section):
+    # transpose it back to the physical shape because Triton looks at the physical storage
+    # note: q and k are incontiguous before the transformation and will become contiguous after transpose
+    q = q.transpose(1, 2)
+    k = k.transpose(1, 2)
+
+    batch_size, seq_len, n_q_head, head_dim = q.shape
+    n_kv_head = k.shape[2]
+    pad_hd = triton.next_power_of_2(head_dim)
+    pad_n_q_head = triton.next_power_of_2(n_q_head)
+    pad_n_kv_head = triton.next_power_of_2(n_kv_head)
+
+    n_row = batch_size * seq_len
+
+    # ensure tensors passed into the kernel are contiguous. It will be no-op if they are already contiguous
+    q = q.contiguous()
+    k = k.contiguous()
+    cos = cos.contiguous()
+    sin = sin.contiguous()
+
+    # Compute tiling strategy based on UB capacity
+    dtype_size = q.element_size()
+    # MROPE forward tiling strategy:
+    # - cos_vals and sin_vals (include text, height and width) are loaded once outside loops (shared): (pad_hd // 2) * 4 = 2 * pad_hd elements each
+    # - In q heads loop (peak memory):
+    #   * q_left: BLOCK_Q * (pad_hd // 2) elements
+    #   * q_right: BLOCK_Q * (pad_hd // 2) elements
+    #   * new_left: BLOCK_Q * (pad_hd // 2) elements (intermediate result)
+    #   * new_right: BLOCK_Q * (pad_hd // 2) elements (intermediate result)
+    #   * Total: 4 * BLOCK_Q * (pad_hd // 2) = 2 * BLOCK_Q * pad_hd elements
+    # - In k heads loop (peak memory):
+    #   * k_left: BLOCK_K * (pad_hd // 2) elements
+    #   * k_right: BLOCK_K * (pad_hd // 2) elements
+    #   * new_left: BLOCK_K * (pad_hd // 2) elements (intermediate result)
+    #   * new_right: BLOCK_K * (pad_hd // 2) elements (intermediate result)
+    #   * Total: 4 * BLOCK_K * (pad_hd // 2) = 2 * BLOCK_K * pad_hd elements
+    # - Since q and k are processed separately, peak memory is max(BLOCK_Q, BLOCK_K) case
+    # - Plus shared cos/sin: 2 * (pad_hd // 2) = pad_hd elements
+    # - Conservative estimate: (2 * BLOCK_SIZE * pad_hd + pad_hd) * dtype_size * 8 bits
+    # - Simplified: (2 * BLOCK_SIZE + 2) * pad_hd * dtype_size * 8 bits
+    # - For safety, use: memory_multiplier=3.0 * BLOCK_SIZE * pad_hd * dtype_size * 8 bits
+    # - shapes: ((pad_n_q_head, pad_hd), (pad_n_kv_head, pad_hd))
+    # - tiling_dims: (0, 0) means first dimension of each shape can be tiled
+    # - Returns: ((block_size_q, pad_hd), (block_size_kv, pad_hd))
+    shapes = ((pad_n_q_head, pad_hd), (pad_n_kv_head, pad_hd))
+    tile_shapes = compute_default_tiling_strategy(
+        safety_margin=0.90,
+        dtype_size=dtype_size,
+        memory_multiplier=3.0,
+        shapes=shapes,
+        tiling_dims=(0, 0),
+    )
+
+    if tile_shapes is not None and len(tile_shapes) == len(shapes):
+        # Strategy returns ((block_size_q, pad_hd), (block_size_kv, pad_hd))
+        q_tile_shape, k_tile_shape = tile_shapes
+        BLOCK_Q, _ = q_tile_shape
+        BLOCK_K, _ = k_tile_shape
+    else:
+        # Fallback to conservative defaults
+        BLOCK_Q = triton.next_power_of_2(pad_n_q_head)
+        BLOCK_K = triton.next_power_of_2(pad_n_kv_head)
+    _triton_qwen2vl_mrope_npu[(n_row,)](
+        q,
+        q.stride(1),
+        k,
+        k.stride(1),
+        cos,
+        sin,
+        seq_len,
+        batch_size,
+        n_q_head,
+        n_kv_head,
+        head_dim,
+        mrope_section[0],
+        mrope_section[1],
+        BLOCK_Q,
+        BLOCK_K,
+        BACKWARD_PASS=False,
+    )
+    return q.transpose(1, 2), k.transpose(1, 2), cos, sin
+
+
+def qwen2vl_mrope_backward(dq, dk, cos, sin, mrope_section):
+    dq = dq.transpose(1, 2)
+    dk = dk.transpose(1, 2)
+
+    batch_size, seq_len, n_q_head, head_dim = dq.shape
+    n_kv_head = dk.shape[2]
+    pad_hd = triton.next_power_of_2(head_dim)
+    pad_n_q_head = triton.next_power_of_2(n_q_head)
+    pad_n_kv_head = triton.next_power_of_2(n_kv_head)
+
+    n_row = batch_size * seq_len
+
+    # ensure dq and dk are contiguous
+    dq = dq.contiguous()
+    dk = dk.contiguous()
+
+    # Compute tiling strategy based on UB capacity
+    dtype_size = dq.element_size()
+    # MROPE backward tiling strategy:
+    # - cos_vals and sin_vals (include text, height and width) are loaded once outside loops (shared): (pad_hd // 2) * 4 = 2 * pad_hd elements each
+    # - In q heads loop (peak memory):
+    #   * q_left: BLOCK_Q * (pad_hd // 2) elements
+    #   * q_right: BLOCK_Q * (pad_hd // 2) elements
+    #   * new_left: BLOCK_Q * (pad_hd // 2) elements (intermediate result)
+    #   * new_right: BLOCK_Q * (pad_hd // 2) elements (intermediate result)
+    #   * Total: 4 * BLOCK_Q * (pad_hd // 2) = 2 * BLOCK_Q * pad_hd elements
+    # - In k heads loop (peak memory):
+    #   * k_left: BLOCK_K * (pad_hd // 2) elements
+    #   * k_right: BLOCK_K * (pad_hd // 2) elements
+    #   * new_left: BLOCK_K * (pad_hd // 2) elements (intermediate result)
+    #   * new_right: BLOCK_K * (pad_hd // 2) elements (intermediate result)
+    #   * Total: 4 * BLOCK_K * (pad_hd // 2) = 2 * BLOCK_K * pad_hd elements
+    # - Since q and k are processed separately, peak memory is max(BLOCK_Q, BLOCK_K) case
+    # - Plus shared cos/sin: 2 * (pad_hd // 2) = pad_hd elements
+    # - Conservative estimate: (2 * BLOCK_SIZE * pad_hd + pad_hd) * dtype_size * 8 bits
+    # - Simplified: (2 * BLOCK_SIZE + 2) * pad_hd * dtype_size * 8 bits
+    # - For safety, use: memory_multiplier=3.0 * BLOCK_SIZE * pad_hd * dtype_size * 8 bits
+    # - shapes: ((pad_n_q_head, pad_hd), (pad_n_kv_head, pad_hd))
+    # - tiling_dims: (0, 0) means first dimension of each shape can be tiled
+    # - Returns: ((block_size_q, pad_hd), (block_size_kv, pad_hd))
+    shapes = ((pad_n_q_head, pad_hd), (pad_n_kv_head, pad_hd))
+    tile_shapes = compute_default_tiling_strategy(
+        safety_margin=0.90,
+        dtype_size=dtype_size,
+        memory_multiplier=3.0,
+        shapes=shapes,
+        tiling_dims=(0, 0),
+    )
+
+    if tile_shapes is not None and len(tile_shapes) == len(shapes):
+        # Strategy returns ((block_size_q, pad_hd), (block_size_kv, pad_hd))
+        q_tile_shape, k_tile_shape = tile_shapes
+        BLOCK_Q, _ = q_tile_shape
+        BLOCK_K, _ = k_tile_shape
+    else:
+        # Fallback to conservative defaults
+        BLOCK_Q = triton.next_power_of_2(pad_n_q_head)
+        BLOCK_K = triton.next_power_of_2(pad_n_kv_head)
+    _triton_qwen2vl_mrope_npu[(n_row,)](
+        dq,
+        dq.stride(1),
+        dk,
+        dk.stride(1),
+        cos,
+        sin,
+        seq_len,
+        batch_size,
+        n_q_head,
+        n_kv_head,
+        head_dim,
+        mrope_section[0],
+        mrope_section[1],
+        BLOCK_Q,
+        BLOCK_K,
+        BACKWARD_PASS=True,
+    )
+    return dq.transpose(1, 2), dk.transpose(1, 2)
+
+
+class LigerQwen2VLMRopeFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, cos, sin, mrope_section, unsqueeze_dim=1):
+        """
+        q size: (bsz, n_q_head, seq_len, head_dim)
+        k size: (bsz, n_kv_head, seq_len, head_dim)
+        cos size: (3, bsz, seq_len, head_dim)
+        sin size: (3, bsz, seq_len, head_dim)
+        """
+        q, k, cos, sin = qwen2vl_mrope_forward(q, k, cos, sin, mrope_section)
+        ctx.save_for_backward(cos, sin)
+        ctx.mrope_section = mrope_section
+        return q, k
+
+    def backward(ctx, dq, dk):
+        """
+        dq size: (bsz, n_q_head, seq_len, head_dim)
+        dk size: (bsz, n_kv_head, seq_len, head_dim)
+        cos size: (3, bsz, seq_len, head_dim)
+        sin size: (3, bsz, seq_len, head_dim)
+        """
+        cos, sin = ctx.saved_tensors
+        mrope_section = ctx.mrope_section
+        dq, dk = qwen2vl_mrope_backward(dq, dk, cos, sin, mrope_section)
+        return dq, dk, None, None, None, None

{liger_kernel_nightly-0.6.4.dev20251229100549 → liger_kernel_nightly-0.6.4.dev20260113013032}/src/liger_kernel/ops/backends/_ascend/ops/rope.py

@@ -239,8 +239,8 @@ def rope_backward(dq, dk, cos, sin):
         BLOCK_K, _ = k_tile_shape
     else:
         # Fallback to conservative defaults
-        BLOCK_Q = min(32, triton.next_power_of_2(pad_n_q_head))
-        BLOCK_K = min(32, triton.next_power_of_2(pad_n_kv_head))
+        BLOCK_Q = triton.next_power_of_2(pad_n_q_head)
+        BLOCK_K = triton.next_power_of_2(pad_n_kv_head)
 
     _triton_rope_npu[(n_row,)](
         dq,