liger-kernel 0.6.3__py3-none-any.whl → 0.6.5__py3-none-any.whl

liger_kernel/ops/backends/_ascend/ops/llama4_rope.py

@@ -0,0 +1,298 @@
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.backends._ascend.ub_manager import compute_default_tiling_strategy
+
+
+def _prepare_freqs(freqs_cis: torch.Tensor, seq_len: int, head_dim_half: int):
+    """
+    Canonicalize freqs to (seq_len, head_dim_half) real/imag tensors.
+
+    Supports:
+    - complex freqs: (..., head_dim_half) complex -> real/imag
+    - packed freqs: (..., 2*head_dim_half) real -> split into real/imag
+    """
+    if freqs_cis.is_complex():
+        freqs_real = freqs_cis.real
+        freqs_imag = freqs_cis.imag
+    else:
+        if freqs_cis.shape[-1] == 2 * head_dim_half:
+            freqs_real = freqs_cis[..., :head_dim_half]
+            freqs_imag = freqs_cis[..., head_dim_half:]
+        else:
+            raise ValueError(
+                f"Unexpected freqs_cis shape for non-complex input: {freqs_cis.shape}, "
+                f"expected last dim = {2 * head_dim_half}"
+            )
+
+    if freqs_real.shape[-1] != head_dim_half:
+        raise ValueError(f"Unexpected last dim for freqs: {freqs_real.shape[-1]} (expected {head_dim_half})")
+
+    # Flatten leading dims -> (N, head_dim_half)
+    freqs_real = freqs_real.reshape(-1, head_dim_half)
+    freqs_imag = freqs_imag.reshape(-1, head_dim_half)
+
+    # Broadcast/slice to (seq_len, head_dim_half)
+    if freqs_real.shape[0] < seq_len:
+        if freqs_real.shape[0] == 1:
+            freqs_real = freqs_real.expand(seq_len, -1)
+            freqs_imag = freqs_imag.expand(seq_len, -1)
+        else:
+            raise ValueError(f"Insufficient rows in freqs: {freqs_real.shape[0]} < seq_len={seq_len}")
+    elif freqs_real.shape[0] > seq_len:
+        freqs_real = freqs_real[:seq_len]
+        freqs_imag = freqs_imag[:seq_len]
+
+    return freqs_real, freqs_imag
+
+
+def _cast_and_contiguous(q, k, freqs_real, freqs_imag):
+    # Align dtype: fp32 only when q is fp32; otherwise keep q dtype for perf
+    compute_dtype = torch.float32 if q.dtype == torch.float32 else q.dtype
+
+    if k.dtype != q.dtype:
+        k = k.to(q.dtype)
+
+    q = q.to(compute_dtype).contiguous()
+    k = k.to(compute_dtype).contiguous()
+    freqs_real = freqs_real.to(compute_dtype).contiguous()
+    freqs_imag = freqs_imag.to(compute_dtype).contiguous()
+    return q, k, freqs_real, freqs_imag, compute_dtype
+
+
+@triton.jit
+def _triton_llama4_rope_npu(
+    q_ptr,
+    k_ptr,
+    freqs_real_ptr,
+    freqs_imag_ptr,
+    q_row_stride,
+    k_row_stride,
+    q_head_stride,
+    k_head_stride,
+    freqs_row_stride,
+    sl,
+    bs: tl.constexpr,
+    n_qh: tl.constexpr,
+    n_kh: tl.constexpr,
+    hd: tl.constexpr,
+    BLOCK_Q: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    imag_sign: tl.constexpr,
+):
+    """
+    Llama4 RoPE on Ascend NPU for interleaved complex layout:
+    - q/k shape: (bs, sl, n_heads, hd)
+    - last dim layout: [real0, imag0, real1, imag1, ...]
+    - freqs_real/imag: (sl, hd//2)
+    """
+    pid = tl.program_id(0).to(tl.int64)
+    batch_idx = pid // sl
+    seq_idx = pid % sl
+
+    if batch_idx >= bs:
+        return
+
+    q_base = q_ptr + pid * q_row_stride
+    k_base = k_ptr + pid * k_row_stride
+
+    freq_base = seq_idx * freqs_row_stride
+    hd_idx = tl.arange(0, hd)
+    hd_mask = hd_idx < (hd)
+
+    freq_idx = tl.arange(0, hd // 2)
+    freq_mask = freq_idx < (hd // 2)
+
+    freqs_real = tl.load(freqs_real_ptr + freq_base + freq_idx, mask=freq_mask, other=0.0)
+    freqs_imag = tl.load(freqs_imag_ptr + freq_base + freq_idx, mask=freq_mask, other=0.0) * imag_sign
+
+    # Q heads (chunked for UB)
+    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] & hd_mask[None, :]
+
+        head_ptr = q_base + qh_idx[:, None] * q_head_stride
+
+        q_pair = tl.load(
+            head_ptr + hd_idx[None, :],
+            mask=block_mask,
+            other=0.0,
+        )
+        q_pair = q_pair.reshape(BLOCK_Q, hd // 2, 2, can_reorder=True)
+        q_real, q_imag = tl.split(q_pair)
+
+        new_real = tl.math.fma(q_real, freqs_real, -(q_imag * freqs_imag))
+        new_imag = tl.math.fma(q_real, freqs_imag, q_imag * freqs_real)
+        new_q_pair = tl.interleave(new_real, new_imag)
+
+        tl.store(head_ptr + hd_idx[None, :], new_q_pair, mask=block_mask)
+
+    # K heads (chunked for UB)
+    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] & hd_mask[None, :]
+
+        head_ptr = k_base + kh_idx[:, None] * k_head_stride
+
+        k_pair = tl.load(
+            head_ptr + hd_idx[None, :],
+            mask=block_mask,
+            other=0.0,
+        )
+
+        k_pair = k_pair.reshape(BLOCK_K, hd // 2, 2, can_reorder=True)
+        k_real, k_imag = tl.split(k_pair)
+
+        new_real = tl.math.fma(k_real, freqs_real, -(k_imag * freqs_imag))
+        new_imag = tl.math.fma(k_real, freqs_imag, k_imag * freqs_real)
+        new_k_pair = tl.interleave(new_real, new_imag)
+
+        tl.store(head_ptr + hd_idx[None, :], new_k_pair, mask=block_mask)
+
+
+def llama4_rope_forward(q, k, freqs_cis):
+    """
+    Ascend NPU implementation of Llama4 RoPE.
+
+    q/k: (bs, sl, n_heads, hd) with interleaved complex last-dim layout.
+    freqs_cis: complex (..., hd//2) OR packed (..., 2*(hd//2)).
+    """
+    original_dtype = q.dtype
+
+    bs, sl, n_qh, hd = q.shape
+    _, _, n_kh, _ = k.shape
+    if hd % 2 != 0:
+        raise ValueError(f"head_dim must be even for interleaved complex layout, got {hd}")
+    hd_half = hd // 2
+
+    freqs_real, freqs_imag = _prepare_freqs(freqs_cis, sl, hd_half)
+    q, k, freqs_real, freqs_imag, compute_dtype = _cast_and_contiguous(q, k, freqs_real, freqs_imag)
+
+    # UB tiling strategy: tile heads dimension only
+    dtype_size = q.element_size()
+    shapes = ((n_qh, hd), (n_kh, hd))
+    tile_shapes = compute_default_tiling_strategy(
+        safety_margin=0.90,
+        dtype_size=dtype_size,
+        memory_multiplier=12.0,
+        shapes=shapes,
+        tiling_dims=(0, 0),
+    )
+
+    if tile_shapes is not None and len(tile_shapes) == len(shapes):
+        q_tile_shape, k_tile_shape = tile_shapes
+        BLOCK_Q, _ = q_tile_shape
+        BLOCK_K, _ = k_tile_shape
+    else:
+        BLOCK_Q = triton.next_power_of_2(n_qh)
+        BLOCK_K = triton.next_power_of_2(n_kh)
+
+    n_row = bs * sl
+
+    _triton_llama4_rope_npu[(n_row,)](
+        q,
+        k,
+        freqs_real,
+        freqs_imag,
+        q.stride(1),
+        k.stride(1),
+        q.stride(2),
+        k.stride(2),
+        freqs_real.stride(0),
+        sl,
+        bs,
+        n_qh,
+        n_kh,
+        hd,
+        BLOCK_Q,
+        BLOCK_K,
+        imag_sign=1.0,
+    )
+
+    if compute_dtype != original_dtype:
+        q = q.to(original_dtype)
+        k = k.to(original_dtype)
+    return q, k
+
+
+def llama4_rope_backward(dq, dk, freqs_cis):
+    """
+    Ascend NPU implementation of Llama4 RoPE.
+
+    q/k: (bs, sl, n_heads, hd) with interleaved complex last-dim layout.
+    freqs_cis: complex (..., hd//2) OR packed (..., 2*(hd//2)).
+    """
+    original_dtype = dq.dtype
+
+    bs, sl, n_qh, hd = dq.shape
+    _, _, n_kh, _ = dk.shape
+    if hd % 2 != 0:
+        raise ValueError(f"head_dim must be even for interleaved complex layout, got {hd}")
+    hd_half = hd // 2
+
+    freqs_real, freqs_imag = _prepare_freqs(freqs_cis, sl, hd_half)
+    dq, dk, freqs_real, freqs_imag, compute_dtype = _cast_and_contiguous(dq, dk, freqs_real, freqs_imag)
+
+    # UB tiling strategy: tile heads dimension only
+    dtype_size = dq.element_size()
+    shapes = ((n_qh, hd), (n_kh, hd))
+    tile_shapes = compute_default_tiling_strategy(
+        safety_margin=0.90,
+        dtype_size=dtype_size,
+        memory_multiplier=12.0,
+        shapes=shapes,
+        tiling_dims=(0, 0),
+    )
+
+    if tile_shapes is not None and len(tile_shapes) == len(shapes):
+        q_tile_shape, k_tile_shape = tile_shapes
+        BLOCK_Q, _ = q_tile_shape
+        BLOCK_K, _ = k_tile_shape
+    else:
+        BLOCK_Q = triton.next_power_of_2(n_qh)
+        BLOCK_K = triton.next_power_of_2(n_kh)
+
+    n_row = bs * sl
+
+    _triton_llama4_rope_npu[(n_row,)](
+        dq,
+        dk,
+        freqs_real,
+        freqs_imag,
+        dq.stride(1),
+        dk.stride(1),
+        dq.stride(2),
+        dk.stride(2),
+        freqs_real.stride(0),
+        sl,
+        bs,
+        n_qh,
+        n_kh,
+        hd,
+        BLOCK_Q,
+        BLOCK_K,
+        imag_sign=-1.0,
+    )
+
+    if compute_dtype != original_dtype:
+        dq = dq.to(original_dtype)
+        dk = dk.to(original_dtype)
+    return dq, dk
+
+
+class LigerLlama4RopeFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, freqs_cis, BLOCK_SIZE: int = None):
+        # BLOCK_SIZE is ignored for Ascend (we auto-tile heads by UB), kept for API compatibility
+        q_out, k_out = llama4_rope_forward(q, k, freqs_cis)
+        ctx.save_for_backward(freqs_cis.detach() if isinstance(freqs_cis, torch.Tensor) else freqs_cis)
+        return q_out, k_out
+
+    @staticmethod
+    def backward(ctx, dq, dk):
+        (freqs_cis,) = ctx.saved_tensors
+        dq_out, dk_out = llama4_rope_backward(dq, dk, freqs_cis)
+        return dq_out, dk_out, None, None

liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py

@@ -0,0 +1,275 @@
+import torch
+import triton
+import triton.language as tl
+
+from liger_kernel.ops.backends._ascend.ub_manager import compute_default_tiling_strategy
+from liger_kernel.ops.utils import get_npu_core_count
+
+
+@triton.jit
+def _triton_qwen2vl_mrope_npu(
+    q_ptr,
+    q_row_stride,
+    k_ptr,
+    k_row_stride,
+    cos,
+    sin,
+    sl,
+    bs: tl.constexpr,
+    total_rows: 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,
+    NUM_STAGES: tl.constexpr,
+    BACKWARD_PASS: tl.constexpr = False,
+):
+    program_id = tl.program_id(0)
+    num_programs = tl.num_programs(0)
+
+    rows_per_program = (total_rows + num_programs - 1) // num_programs
+    start_row = program_id * rows_per_program
+    actual_rows = tl.minimum(rows_per_program, total_rows - start_row)
+
+    for row_offset in tl.range(0, actual_rows, num_stages=NUM_STAGES):
+        pid = start_row + row_offset
+
+        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))
+
+        # Process q heads in chunks to prevent UB overflow
+        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)
+
+        # Process k heads in chunks to prevent UB overflow
+        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 get_optimal_block_size_mrope(pad_n_q_head, pad_n_kv_head, pad_hd, dtype_size):
+    # MROPE forward tiling strategy:
+    # - cos_vals and sin_vals (include text, height and width) are loaded once outside loops (shared): (pad_hd // 2) * 6 = 3 * 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: 6 * (pad_hd // 2) = 3 * pad_hd elements
+    # - Conservative estimate: (2 * BLOCK_SIZE * pad_hd + 3 * pad_hd) * dtype_size * 8 bits
+    # - Simplified: (2 * BLOCK_SIZE + 3) * 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 = 2048
+        BLOCK_K = 2048
+
+    return BLOCK_Q, BLOCK_K
+
+
+def qwen2vl_mrope_forward(q, k, cos, sin, mrope_section):
+    # transpose it back to the physical shape because Triton looks at the physical storage
+    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
+    q = q.contiguous()
+    k = k.contiguous()
+    cos = cos.contiguous()
+    sin = sin.contiguous()
+
+    dtype_size = q.element_size()
+    BLOCK_Q, BLOCK_K = get_optimal_block_size_mrope(pad_n_q_head, pad_n_kv_head, pad_hd, dtype_size)
+
+    num_cores = get_npu_core_count()
+    grid_size = min(num_cores, n_row)
+
+    _triton_qwen2vl_mrope_npu[(grid_size,)](
+        q,
+        q.stride(1),
+        k,
+        k.stride(1),
+        cos,
+        sin,
+        seq_len,
+        batch_size,
+        n_row,
+        n_q_head,
+        n_kv_head,
+        head_dim,
+        mrope_section[0],
+        mrope_section[1],
+        BLOCK_Q,
+        BLOCK_K,
+        NUM_STAGES=3,
+        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()
+
+    dtype_size = dq.element_size()
+    BLOCK_Q, BLOCK_K = get_optimal_block_size_mrope(pad_n_q_head, pad_n_kv_head, pad_hd, dtype_size)
+
+    num_cores = get_npu_core_count()
+    grid_size = min(num_cores, n_row)
+
+    _triton_qwen2vl_mrope_npu[(grid_size,)](
+        dq,
+        dq.stride(1),
+        dk,
+        dk.stride(1),
+        cos,
+        sin,
+        seq_len,
+        batch_size,
+        n_row,
+        n_q_head,
+        n_kv_head,
+        head_dim,
+        mrope_section[0],
+        mrope_section[1],
+        BLOCK_Q,
+        BLOCK_K,
+        NUM_STAGES=3,
+        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
+
+    @staticmethod
+    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