init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/ops/triton/fla/chunk.py

@@ -0,0 +1,391 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+import warnings
+
+import torch
+from einops import rearrange
+from vllm.distributed import get_pcp_group
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.fla.ops.utils import SUPPRESS_LEVEL
+
+from vllm_ascend.ops.gdn_attn_builder import _compact_empty_segments
+
+from .chunk_delta_h import chunk_gated_delta_rule_fwd_h  # noqa: F401
+from .chunk_delta_hupdate import chunk_gated_delta_rule_fwd_hupdate
+from .chunk_o import chunk_fwd_o  # noqa: F401
+from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
+from .cumsum import chunk_local_cumsum
+from .l2norm import l2norm_fwd
+from .solve_tril import solve_tril
+from .utils import input_guard, prepare_final_chunk_indices
+from .wy_fast import recompute_w_u_fwd
+
+
+def chunk_gated_delta_rule_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: torch.LongTensor | None = None,
+    prebuilt_meta=None,
+):
+    forward_context = get_forward_context()
+    num_decodes = 0
+    attn_metadata = forward_context.attn_metadata
+    if attn_metadata is not None and isinstance(attn_metadata, dict):
+        attn_metadata = next(iter(attn_metadata.values()), None)
+    if attn_metadata is not None:
+        num_decodes = attn_metadata.num_decodes
+    chunk_size = 64
+    block_indices_cumsum = None if prebuilt_meta is None else prebuilt_meta.block_indices_cumsum
+    cu_seqlens_host = None if prebuilt_meta is None else prebuilt_meta.cu_seqlens_host
+    chunk_indices_chunk64 = None if prebuilt_meta is None else prebuilt_meta.chunk_indices_chunk64
+    chunk_indices_chunk64_host = None if prebuilt_meta is None else prebuilt_meta.chunk_indices_chunk64_host
+    chunk_offsets_chunk64 = None if prebuilt_meta is None else prebuilt_meta.chunk_offsets_chunk64
+    update_chunk_offsets_chunk64 = None if prebuilt_meta is None else prebuilt_meta.update_chunk_offsets_chunk64
+    final_chunk_indices_chunk64 = None if prebuilt_meta is None else prebuilt_meta.final_chunk_indices_chunk64
+    chunk_indices_large_block = None if prebuilt_meta is None else prebuilt_meta.chunk_indices_large_block
+    g = chunk_local_cumsum(
+        g,
+        chunk_size=chunk_size,
+        cu_seqlens=cu_seqlens,
+        block_indices=block_indices_cumsum,
+    )
+    # obtain WY representation. u is actually the new v.
+    A = chunk_scaled_dot_kkt_fwd(
+        k=k,
+        beta=beta,
+        g_cumsum=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices_chunk64,
+        output_dtype=torch.float32,
+    )
+    A = solve_tril(
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices_large_block=chunk_indices_large_block,
+        chunk_indices_bt=chunk_indices_chunk64,
+        output_dtype=k.dtype,
+    )
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g_cumsum=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices_chunk64,
+    )
+
+    k_ascendc = k.to(torch.bfloat16).transpose(1, 2).contiguous()
+    w_ascendc = w.to(torch.bfloat16).transpose(1, 2).contiguous()
+    u_ascendc = u.to(torch.bfloat16).transpose(1, 2).contiguous()
+    g_ascendc = g.transpose(1, 2).contiguous()
+    q_ascendc = q.to(torch.bfloat16).transpose(1, 2).contiguous()
+
+    cu_seqlens = None if cu_seqlens is None else cu_seqlens.to(torch.int64)
+    chunk_indices = None if chunk_indices_chunk64 is None else chunk_indices_chunk64.to(torch.int64)
+    if cu_seqlens_host is None and cu_seqlens is not None:
+        cu_seqlens_host = tuple(cu_seqlens.tolist())
+    if chunk_indices_chunk64_host is None and chunk_indices is not None:
+        chunk_indices_chunk64_host = tuple(chunk_indices.flatten().tolist())
+    # Compact zero-length segments for the AscendC kernels (see
+    # _compact_empty_segments).  chunk_indices_chunk64 is already compact-
+    # ranked and is reused as-is; only cu_seqlens / initial_state need
+    # compacting.
+    if prebuilt_meta is not None and hasattr(prebuilt_meta, "keep_meta"):
+        cu_seqlens_kern = cu_seqlens_host if prebuilt_meta.cu_seqlens_kern is None else prebuilt_meta.cu_seqlens_kern
+        keep_meta = prebuilt_meta.keep_meta
+        initial_state_kern = (
+            initial_state[keep_meta] if initial_state is not None and keep_meta is not None else initial_state
+        )
+    else:
+        cu_seqlens_kern, initial_state_kern, keep_meta = _compact_empty_segments(
+            cu_seqlens_host,
+            initial_state,
+            device=initial_state.device if initial_state is not None else None,
+        )
+    h, v_new, final_state = torch.ops._C_ascend.chunk_gated_delta_rule_fwd_h(
+        k_ascendc,
+        w_ascendc,
+        u_ascendc,
+        g=g_ascendc,
+        gk=None,
+        initial_state=initial_state_kern,
+        output_final_state=True,
+        chunk_size=64,
+        save_new_value=True,
+        cu_seqlens=cu_seqlens_kern,
+        chunk_indices=chunk_indices_chunk64_host,
+        use_exp2=False,
+        transpose_state_layout=False,
+    )
+    if keep_meta is not None:
+        # Scatter the compacted final_state back to the original [N, H, K, V]
+        # layout the PCP state recursion expects; empty segments keep their
+        # initial state.
+        _fs_full = initial_state.clone()
+        _fs_full[keep_meta] = final_state
+        final_state = _fs_full
+
+    if get_pcp_group().world_size > 1:
+        # When integrating mtp, since `mix_qkv` has been split, `num_decode`
+        # cannot be directly obtained from the metadata and needs to be recalculated.
+        actual_num_decodes = getattr(prebuilt_meta, "num_decodes", None)
+        if actual_num_decodes is None:
+            actual_num_decodes = num_decodes
+        h_update = chunk_gated_delta_rule_fwd_hupdate(
+            k=k,
+            w=w,
+            u=u,
+            g=g,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices_chunk64,
+            chunk_offsets=chunk_offsets_chunk64,
+            update_chunk_offsets=update_chunk_offsets_chunk64,
+            num_decodes=actual_num_decodes,
+        )
+        all_final_state = get_pcp_group().all_gather(final_state.unsqueeze(0), 0)
+        final_chunk_indices = final_chunk_indices_chunk64
+        if final_chunk_indices is None:
+            final_chunk_indices = prepare_final_chunk_indices(cu_seqlens, chunk_size)
+        final_h_update = h_update[:, final_chunk_indices, :, :, :]
+        all_final_h_update = get_pcp_group().all_gather(final_h_update, 0)
+
+        updated_state = final_state.new_empty(get_pcp_group().world_size, *final_state.shape)
+        updated_state[0, ...] = all_final_state[0]
+        for i in range(1, get_pcp_group().world_size):
+            # correct_i = all_final_state[i] + Phi_i * (correct_{i-1} - s0)
+            updated_final_state = all_final_state[i] + torch.matmul(
+                all_final_h_update[i, ...], updated_state[i - 1, ...] - initial_state
+            )
+            updated_state[i, ...] = updated_final_state
+
+        final_state = updated_state[-1, ...]
+
+        if get_pcp_group().rank_in_group == 0:
+            updated_h_state = torch.zeros_like(final_state)
+        else:
+            updated_h_state = updated_state[get_pcp_group().rank_in_group - 1, ...]
+
+        if get_pcp_group().rank_in_group > 0:
+            rerun_initial_state = initial_state.clone()
+            prefill_seq_offset = actual_num_decodes
+            prefill_slice = slice(prefill_seq_offset, final_state.shape[0])
+            rerun_initial_state[prefill_slice] = updated_h_state[prefill_slice]
+            h, v_new, _ = chunk_gated_delta_rule_fwd_h(
+                k=k,
+                w=w,
+                u=u,
+                g=g,
+                initial_state=rerun_initial_state,
+                output_final_state=True,
+                cu_seqlens=cu_seqlens,
+                chunk_indices=chunk_indices_chunk64,
+                chunk_offsets=chunk_offsets_chunk64,
+            )
+            h = h.transpose(1, 2).contiguous()
+            v_new = v_new.transpose(1, 2).contiguous()
+
+    o_ascendc = torch.ops._C_ascend.chunk_fwd_o(
+        q_ascendc,
+        k_ascendc,
+        v_new,
+        h,
+        scale,
+        g=g_ascendc,
+        g_gamma=None,
+        cu_seqlens=cu_seqlens_host,
+        chunk_indices=chunk_indices_chunk64_host,
+        chunk_size=64,
+        transpose_state_layout=False,
+    )
+
+    o = o_ascendc.to(torch.bfloat16).transpose(1, 2).contiguous()
+    v_new = v_new.to(torch.bfloat16).transpose(1, 2).contiguous()
+    h = h.to(torch.bfloat16).transpose(1, 2).contiguous()
+
+    if SUPPRESS_LEVEL < 3:
+        return g, o, A, final_state, None, None, None
+    elif SUPPRESS_LEVEL >= 3:
+        return g, o, A, final_state, w, h, v_new
+
+
+class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
+    @staticmethod
+    @input_guard
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        scale: float,
+        initial_state: torch.Tensor,
+        output_final_state: bool,
+        cu_seqlens: torch.LongTensor | None = None,
+        prebuilt_meta=None,
+        use_qk_l2norm_in_kernel: bool = False,
+    ):
+        if use_qk_l2norm_in_kernel:
+            q = l2norm_fwd(q)
+            k = l2norm_fwd(k)
+        g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            scale=scale,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            cu_seqlens=cu_seqlens,
+            prebuilt_meta=prebuilt_meta,
+        )
+        ctx.scale = scale
+        ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
+        return o.to(q.dtype), final_state
+
+
+@torch.compiler.disable
+def chunk_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    cu_seqlens: torch.LongTensor | None = None,
+    prebuilt_meta=None,
+    head_first: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
+    core_attn_out: torch.Tensor | None = None,
+):
+    r"""
+    Args:
+        q (torch.Tensor):
+            queries of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        k (torch.Tensor):
+            keys of shape `[B, T, H, K]` if `head_first=False` else `[B, H, T, K]`.
+        v (torch.Tensor):
+            values of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
+        g (torch.Tensor):
+            (forget) gating tensor (in log space!) of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
+        beta (torch.Tensor):
+            betas of shape `[B, T, H]` if `head_first=False` else `[B, H, T]`.
+        scale (Optional[int]):
+            Scale factor for the RetNet attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `[N, H, K, V]` for `N` input sequences.
+            For equal-length input sequences, `N` equals the batch size `B`.
+            Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
+        cu_seqlens (torch.LongTensor):
+            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
+            consistent with the FlashAttention API.
+        head_first (Optional[bool]):
+            Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
+            Default: `False`.
+
+    Returns:
+        o (torch.Tensor):
+            Outputs of shape `[B, T, H, V]` if `head_first=False` else `[B, H, T, V]`.
+        final_state (torch.Tensor):
+            Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
+
+    Examples::
+        >>> import torch
+        >>> import torch.nn.functional as F
+        >>> from einops import rearrange
+        >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+        # inputs with equal lengths
+        >>> B, T, H, K, V = 4, 2048, 4, 512, 512
+        >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
+        >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
+        >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
+        >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
+        >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
+        >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')
+        >>> o, ht = chunk_gated_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True
+        )
+        # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
+        >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
+        # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
+        >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
+        >>> o_var, ht_var = chunk_gated_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True,
+            cu_seqlens=cu_seqlens
+        )
+    """
+    assert q.dtype == k.dtype == v.dtype
+    assert q.dtype != torch.float32, "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
+    assert len(beta.shape) == 3, "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
+
+    if head_first:
+        raise DeprecationWarning(
+            "chunk_gated_delta_rule: head_first is deprecated and will be removed in a future version. "
+            "Please use head_first=False for now instead.",
+            stacklevel=2,
+        )
+        q, k, v, beta, g = map(lambda x: rearrange(x, "b h t ... -> b t h ..."), (q, k, v, beta, g))
+    if not head_first and q.shape[1] < q.shape[2]:
+        warnings.warn(
+            f"chunk_gated_delta_rule: Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
+            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
+            "when head_first=False was specified. "
+            "Please verify your input tensor format matches the expected shape [B, T, H, ...].",
+            stacklevel=2,
+        )
+    if cu_seqlens is not None:
+        if q.shape[0] != 1:
+            raise ValueError(
+                f"chunk_gated_delta_rule: The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
+                f"Please flatten variable-length inputs before processing."
+            )
+        if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
+            raise ValueError(
+                f"chunk_gated_delta_rule: The number of initial states is expected to be equal to the number of input sequences, "
+                f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."
+            )
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    o, final_state = ChunkGatedDeltaRuleFunction.apply(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        scale,
+        initial_state,
+        output_final_state,
+        cu_seqlens,
+        prebuilt_meta,
+        use_qk_l2norm_in_kernel,
+    )
+    if head_first:
+        o = rearrange(o, "b t h ... -> b h t ...")
+    return o, final_state

vllm_ascend/ops/triton/fla/chunk_delta_h.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .utils import prepare_chunk_indices, prepare_chunk_offsets, safe_exp
+
+_CONDITIONS = ("seq7168",)
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
+        "STORE_FINAL_STATE": lambda args: args["ht"] is not None,
+        "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T", "H", "Hg", "K", "V"])
+def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
+    k,
+    v,
+    w,
+    v_new,
+    g,
+    h,
+    h0,
+    ht,
+    cu_seqlens,
+    chunk_offsets,
+    h_update,
+    T,
+    H,
+    Hg,
+    K,
+    V,
+    BT: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    STORE_FINAL_STATE: tl.constexpr,
+    SAVE_NEW_VALUE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_nh = tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    T_max = 1 * T
+    if IS_VARLEN:
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    stride_v = H * V
+    stride_k = Hg * K
+    stride_w = H * K
+
+    b_h1_bv1 = tl.zeros([128, 64], dtype=tl.float32)
+    b_h1_bv2 = tl.zeros([128, 64], dtype=tl.float32)
+    # create b_hupd_bv1 and b_hupd_bv2
+
+    v_start1 = 0
+    v_start2 = 64
+
+    offs_k = tl.arange(0, 128)[:, None]
+    offs_v1 = v_start1 + tl.arange(0, 64)[None, :]
+    offs_v2 = v_start2 + tl.arange(0, 64)[None, :]
+    mask_kv1 = (offs_k < K) & (offs_v1 < V)
+    mask_kv2 = (offs_k < K) & (offs_v2 < V)
+
+    # load initial state
+    if USE_INITIAL_STATE:
+        h0_ptr = h0 + i_nh * K * V
+        ptr_h0_bv1 = h0_ptr + offs_k * V + offs_v1 * 1
+        b_h1_bv1 += tl.load(ptr_h0_bv1, mask=mask_kv1, other=0.0).to(tl.float32)
+
+        ptr_h0_bv2 = h0_ptr + offs_k * V + offs_v2 * 1
+        b_h1_bv2 += tl.load(ptr_h0_bv2, mask=mask_kv2, other=0.0).to(tl.float32)
+
+    # main recurrence
+    for i_t in range(NT):
+        h_base = h + (boh + i_t) * H * K * V + i_h * K * V
+
+        p_h1_bv1 = tl.make_block_ptr(h_base, (K, V), (V, 1), (0, v_start1), (128, 64), (1, 0))
+        tl.store(p_h1_bv1, b_h1_bv1.to(p_h1_bv1.dtype.element_ty), boundary_check=(0, 1))
+
+        p_h1_bv2 = tl.make_block_ptr(h_base, (K, V), (V, 1), (0, v_start2), (128, 64), (1, 0))
+        tl.store(p_h1_bv2, b_h1_bv2.to(p_h1_bv2.dtype.element_ty), boundary_check=(0, 1))
+
+        offs_t_wv = (i_t * BT + tl.arange(0, BT))[:, None]
+        offs_k_wv = tl.arange(0, 128)[None, :]
+        mask_w = (offs_t_wv < T) & (offs_k_wv < K)
+
+        w_base = w + bos * H * K + i_h * K
+        ptr_w = w_base + offs_t_wv * stride_w + offs_k_wv * 1
+        b_w = tl.load(ptr_w, mask=mask_w, other=0.0)
+
+        k_base = k + bos * Hg * K + (i_h // (H // Hg)) * K
+        p_k = tl.make_block_ptr(k_base, (K, T), (1, stride_k), (0, i_t * BT), (128, BT), (0, 1))
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+
+        v_new_base = v_new + bos * H * V + i_h * V
+
+        last_idx = min((i_t + 1) * BT, T) - 1
+        b_g_last = tl.load(g + bos + i_h * T_max + last_idx)
+
+        offs_t = i_t * BT + tl.arange(0, BT)
+        mask_t = offs_t < T
+        g_ptr = g + bos + i_h * T_max
+        b_g = tl.load(g_ptr + offs_t, mask=mask_t, other=0.0)
+
+        b_g = safe_exp(b_g_last - b_g)
+        b_g_last = tl.exp(b_g_last)
+
+        offs_t_v = (i_t * BT + tl.arange(0, BT))[:, None]
+        mask_v1 = (offs_t_v < T) & (offs_v1 < V)
+
+        v_base = v + bos * H * V + i_h * V
+        ptr_v1 = v_base + offs_t_v * stride_v + offs_v1 * 1
+        b_v1 = tl.load(ptr_v1, mask=mask_v1, other=0.0)
+        b_v_new1 = b_v1.to(tl.float32)
+        b_v_new1 -= tl.dot(b_w, b_h1_bv1.to(b_w.dtype))
+
+        if SAVE_NEW_VALUE:
+            p_v_new1 = tl.make_block_ptr(v_new_base, (T, V), (stride_v, 1), (i_t * BT, v_start1), (BT, 64), (1, 0))
+            tl.store(p_v_new1, b_v_new1.to(p_v_new1.dtype.element_ty), boundary_check=(0, 1))
+
+        if USE_G:
+            b_v_new1 = b_v_new1 * b_g[:, None]
+            b_h1_bv1 = b_h1_bv1 * b_g_last
+
+        b_v_new1 = b_v_new1.to(k.dtype.element_ty)
+        b_h1_bv1 += tl.dot(b_k, b_v_new1)
+
+        mask_v2 = (offs_t_v < T) & (offs_v2 < V)
+        ptr_v2 = v_base + offs_t_v * stride_v + offs_v2 * 1
+        b_v2 = tl.load(ptr_v2, mask=mask_v2, other=0.0)
+        b_v_new2 = b_v2.to(tl.float32)
+        b_v_new2 -= tl.dot(b_w, b_h1_bv2.to(b_w.dtype))
+
+        if SAVE_NEW_VALUE:
+            p_v_new2 = tl.make_block_ptr(v_new_base, (T, V), (stride_v, 1), (i_t * BT, v_start2), (BT, 64), (1, 0))
+            tl.store(p_v_new2, b_v_new2.to(p_v_new2.dtype.element_ty), boundary_check=(0, 1))
+
+        if USE_G:
+            b_v_new2 = b_v_new2 * b_g[:, None]
+            b_h1_bv2 = b_h1_bv2 * b_g_last
+
+        b_v_new2 = b_v_new2.to(k.dtype.element_ty)
+        b_h1_bv2 += tl.dot(b_k, b_v_new2)
+
+    # epilogue
+    if STORE_FINAL_STATE:
+        ht_ptr = ht + i_nh * K * V
+
+        p_ht1_bv1 = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, v_start1), (128, 64), (1, 0))
+        tl.store(p_ht1_bv1, b_h1_bv1.to(p_ht1_bv1.dtype.element_ty), boundary_check=(0, 1))
+
+        p_ht1_bv2 = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, v_start2), (128, 64), (1, 0))
+        tl.store(p_ht1_bv2, b_h1_bv2.to(p_ht1_bv2.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_fwd_h(
+    k: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    g: torch.Tensor | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = False,
+    chunk_size: int = 64,  # SY: remove this argument and force chunk size 64?
+    save_new_value: bool = True,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # This kernel is slightly different from fla to support Q/K with different head numbers.
+    # In fla, Q/K always have the same head number, so Hg is always equal to H.
+    B, T, Hg, K, V = *k.shape, u.shape[-1]
+    H = u.shape[-2]
+    BT = chunk_size
+
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    if cu_seqlens is None:
+        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    else:
+        if chunk_offsets is None:
+            chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
+        N, NT, chunk_offsets = (
+            len(cu_seqlens) - 1,
+            len(chunk_indices),
+            chunk_offsets,
+        )
+    assert K <= 256, "current kernel does not support head dimension larger than 256."
+
+    h = k.new_empty(B, NT, H, K, V)
+    h_update = k.new_empty(B, NT, H, K, K)
+    final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None
+
+    v_new = torch.empty_like(u) if save_new_value else None
+    g = g.transpose(1, 2).contiguous()
+
+    def grid(meta):
+        return (1, N * H)
+
+    chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid](
+        k=k,
+        v=u,
+        w=w,
+        v_new=v_new,
+        g=g,
+        h=h,
+        h0=initial_state,
+        ht=final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        h_update=h_update,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        V=V,
+        BT=BT,
+        num_warps=4,
+        num_stages=2,
+    )
+    return h, v_new, final_state

vllm_ascend/ops/triton/fla/chunk_delta_hupdate.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_update_chunk_offsets, safe_exp
+
+_CONDITIONS = ("seq7168",)
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gated_delta_rule_fwd_kernel_hupdate_blockdim64(
+    k,
+    w,
+    g,
+    cu_seqlens,
+    chunk_offsets,
+    h_update,
+    T,
+    H: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_nh = tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    T_max = 1 * T
+    bos, eos = (
+        tl.load(cu_seqlens + i_n).to(tl.int32),
+        tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+    )
+    T = eos - bos
+    NT = tl.cdiv(T, BT)
+    boh = tl.load(chunk_offsets + i_n).to(tl.int32)
+    if IS_VARLEN:
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    stride_k = Hg * K
+    stride_w = H * K
+
+    # create b_hupd_bv1 and b_hupd_bv2
+    off_hupd_1_top = tl.arange(0, 64)[:, None]
+    off_hupd_2_top = tl.arange(0, 64)[None, :]
+
+    # main recurrence
+    for i_t in range(NT):
+        last_idx = min((i_t + 1) * BT, T) - 1
+        b_g_last = tl.load(g + bos + i_h * T_max + last_idx)
+
+        offs_t = i_t * BT + tl.arange(0, BT)
+        mask_t = offs_t < T
+        g_ptr = g + bos + i_h * T_max
+        b_g = tl.load(g_ptr + offs_t, mask=mask_t, other=0.0)
+
+        b_g = safe_exp(b_g_last - b_g)
+        b_g_last = tl.exp(b_g_last)
+
+        offs_t_wv = (i_t * BT + tl.arange(0, BT))[:, None]
+        w_base = w + bos * H * K + i_h * K
+        # get column-sliced w [BT, 64]
+        offs_w_upd1 = tl.arange(0, 64)[None, :]
+        mask_w_upd1 = (offs_t_wv < T) & (offs_w_upd1 < K)
+        ptr_w_upd1 = w_base + offs_t_wv * stride_w + offs_w_upd1 * 1
+        b_w_upd1 = tl.load(ptr_w_upd1, mask=mask_w_upd1, other=0.0).to(tl.float32)
+
+        offs_w_upd2 = 64 + tl.arange(0, 64)[None, :]
+        mask_w_upd2 = (offs_t_wv < T) & (offs_w_upd2 < K)
+        ptr_w_upd2 = w_base + offs_t_wv * stride_w + offs_w_upd2 * 1
+        b_w_upd2 = tl.load(ptr_w_upd2, mask=mask_w_upd2, other=0.0).to(tl.float32)
+
+        k_base = k + bos * Hg * K + (i_h // (H // Hg)) * K
+        # get row-sliced k [64, T]
+        p_k_upd1 = tl.make_block_ptr(k_base, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1))
+        b_k_upd1 = tl.load(p_k_upd1, boundary_check=(0, 1))
+        p_k_upd2 = tl.make_block_ptr(k_base, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1))
+        b_k_upd2 = tl.load(p_k_upd2, boundary_check=(0, 1))
+
+        if USE_G:
+            b_w_upd1 = b_w_upd1 * b_g[:, None]
+            b_w_upd2 = b_w_upd2 * b_g[:, None]
+
+        # compute [64, BT] @ [BT, 64]
+        b_hupd_local_11 = (off_hupd_1_top == off_hupd_2_top).to(tl.float32)
+        b_hupd_local_22 = (off_hupd_1_top == off_hupd_2_top).to(tl.float32)
+
+        # fp32
+        if USE_G:
+            b_hupd_local_11 = b_hupd_local_11 * b_g_last
+            b_hupd_local_22 = b_hupd_local_22 * b_g_last
+
+        b_hupd_local_11 -= tl.dot(b_k_upd1, b_w_upd1.to(b_k_upd1.dtype))
+        b_hupd_local_22 -= tl.dot(b_k_upd2, b_w_upd2.to(b_k_upd2.dtype))
+        b_hupd_local_12 = -tl.dot(b_k_upd1, b_w_upd2.to(b_k_upd1.dtype)).to(tl.float32)
+        b_hupd_local_21 = -tl.dot(b_k_upd2, b_w_upd1.to(b_k_upd2.dtype)).to(tl.float32)
+
+        hupd_base = h_update + (boh + i_t + i_n) * H * K * K + i_h * K * K
+        p_hupd_11 = tl.make_block_ptr(hupd_base, (K, K), (K, 1), (0, 0), (64, 64), (1, 0))
+        b_hupd_11 = tl.load(p_hupd_11, boundary_check=(1, 0))
+        p_hupd_21 = tl.make_block_ptr(hupd_base, (K, K), (K, 1), (64, 0), (64, 64), (1, 0))
+        b_hupd_21 = tl.load(p_hupd_21, boundary_check=(1, 0))
+        p_hupd_12 = tl.make_block_ptr(hupd_base, (K, K), (K, 1), (0, 64), (64, 64), (1, 0))
+        b_hupd_12 = tl.load(p_hupd_12, boundary_check=(1, 0))
+        p_hupd_22 = tl.make_block_ptr(hupd_base, (K, K), (K, 1), (64, 64), (64, 64), (1, 0))
+        b_hupd_22 = tl.load(p_hupd_22, boundary_check=(1, 0))
+
+        b_hupd11_new = tl.dot(b_hupd_local_11.to(b_hupd_11.dtype), b_hupd_11).to(tl.float32)
+        b_hupd11_new += tl.dot(b_hupd_local_12.to(b_hupd_21.dtype), b_hupd_21)
+
+        b_hupd21_new = tl.dot(b_hupd_local_21.to(b_hupd_11.dtype), b_hupd_11).to(tl.float32)
+        b_hupd21_new += tl.dot(b_hupd_local_22.to(b_hupd_21.dtype), b_hupd_21)
+
+        b_hupd12_new = tl.dot(b_hupd_local_11.to(b_hupd_12.dtype), b_hupd_12).to(tl.float32)
+        b_hupd12_new += tl.dot(b_hupd_local_12.to(b_hupd_22.dtype), b_hupd_22)
+
+        b_hupd22_new = tl.dot(b_hupd_local_21.to(b_hupd_12.dtype), b_hupd_12).to(tl.float32)
+        b_hupd22_new += tl.dot(b_hupd_local_22.to(b_hupd_22.dtype), b_hupd_22)
+
+        hupd_next = h_update + (boh + i_t + i_n + 1) * H * K * K + i_h * K * K
+        p_hupd_11 = tl.make_block_ptr(hupd_next, (K, K), (K, 1), (0, 0), (64, 64), (1, 0))
+        tl.store(p_hupd_11, b_hupd11_new.to(p_hupd_11.dtype.element_ty), boundary_check=(0, 1))
+
+        p_hupd_21 = tl.make_block_ptr(hupd_next, (K, K), (K, 1), (64, 0), (64, 64), (1, 0))
+        tl.store(p_hupd_21, b_hupd21_new.to(p_hupd_21.dtype.element_ty), boundary_check=(0, 1))
+
+        p_hupd_12 = tl.make_block_ptr(hupd_next, (K, K), (K, 1), (0, 64), (64, 64), (1, 0))
+        tl.store(p_hupd_12, b_hupd12_new.to(p_hupd_12.dtype.element_ty), boundary_check=(0, 1))
+
+        p_hupd_22 = tl.make_block_ptr(hupd_next, (K, K), (K, 1), (64, 64), (64, 64), (1, 0))
+        tl.store(p_hupd_22, b_hupd22_new.to(p_hupd_22.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_fwd_hupdate(
+    k: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    g: torch.Tensor | None = None,
+    chunk_size: int = 64,  # SY: remove this argument and force chunk size 64?
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_offsets: torch.Tensor | None = None,
+    update_chunk_offsets: torch.Tensor | None = None,
+    num_decodes: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    # This kernel is slightly different from fla to support Q/K with different head numbers.
+    # In fla, Q/K always have the same head number, so Hg is always equal to H.
+    B, T, Hg, K, _ = *k.shape, u.shape[-1]
+    H = u.shape[-2]
+    BT = chunk_size
+
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    if cu_seqlens is None:
+        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    else:
+        if chunk_offsets is None:
+            chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
+        N, NT, chunk_offsets = (
+            len(cu_seqlens) - 1,
+            len(chunk_indices),
+            chunk_offsets,
+        )
+    assert K <= 256, "current kernel does not support head dimension larger than 256."
+
+    h_update = k.new_empty(B, NT + N, H, K, K, dtype=torch.float32)
+    if cu_seqlens is not None and update_chunk_offsets is None:
+        update_chunk_offsets = prepare_update_chunk_offsets(cu_seqlens, BT)
+    update_indices = update_chunk_offsets[:-1]
+    h_update[:, update_indices, :, :, :] = torch.eye(K, dtype=h_update.dtype, device=h_update.device)
+
+    g = g.transpose(1, 2).contiguous()
+
+    def grid(meta):
+        return (1, N * H)
+
+    chunk_gated_delta_rule_fwd_kernel_hupdate_blockdim64[grid](
+        k=k,
+        w=w,
+        g=g,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        h_update=h_update,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        BT=BT,
+        num_warps=4,
+        num_stages=2,
+    )
+    h_update[:, : num_decodes * 2, :, :, :] = torch.zeros((K, K), dtype=h_update.dtype, device=h_update.device)
+    return h_update

vllm_ascend/ops/triton/fla/chunk_o.py

@@ -0,0 +1,163 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .utils import prepare_chunk_offsets, safe_exp
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["chunk_offsets", "scale", "T", "H", "Hg", "K", "V"])
+def chunk_fwd_kernel_o(
+    q,
+    k,
+    v,
+    h,
+    g,
+    o,
+    cu_seqlens,
+    chunk_offsets,
+    scale,
+    T,
+    H,
+    Hg,
+    K,
+    V,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    T_max = T
+
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int64)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    # offset calculation
+    q += (bos * Hg + i_h // (H // Hg)) * K
+    k += (bos * Hg + i_h // (H // Hg)) * K
+    v += (bos * H + i_h) * V
+    o += (bos * H + i_h) * V
+
+    for i_t in range(NT):
+        i_tg = boh + i_t
+        h_base = h + (i_tg * H + i_h).to(tl.int64) * K * V
+        b_o = tl.zeros([BT, BV], dtype=tl.float32)
+        b_A = tl.zeros([BT, BT], dtype=tl.float32)
+
+        for i_k in range(tl.cdiv(K, BK)):
+            p_q = tl.make_block_ptr(q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+            p_k = tl.make_block_ptr(k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+            p_h = tl.make_block_ptr(h_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+            # [BT, BK]
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            # [BK, BT]
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            # [BK, BV]
+            b_h = tl.load(p_h, boundary_check=(0, 1))
+
+            # [BT, BK] @ [BK, BV] -> [BT, BV]
+            b_o += tl.dot(b_q, b_h)
+            # [BT, BK] @ [BK, BT] -> [BT, BT]
+            b_A += tl.dot(b_q, b_k)
+
+        if USE_G:
+            offs_t = i_t * BT + tl.arange(0, BT)
+            mask_t = offs_t < T
+            g_ptr = g + bos + i_h * T_max
+            b_g = tl.load(g_ptr + offs_t, mask=mask_t, other=0.0)
+
+            b_o = b_o * tl.exp(b_g)[:, None]
+            b_A = b_A * safe_exp(b_g[:, None] - b_g[None, :])
+
+        o_i = tl.arange(0, BT).to(tl.float32)
+        m_A = o_i[:, None] >= o_i[None, :]
+        b_A = tl.where(m_A, b_A, 0)
+
+        p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_o = tl.make_block_ptr(o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+
+        b_v = tl.load(p_v, boundary_check=(0, 1))
+        # to fix mma -> mma layout conversion
+        # already solved by fla v3.2 or higher
+        b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_fwd_o(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    h: torch.Tensor,
+    g: torch.Tensor | None = None,
+    scale: float | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    chunk_offsets: torch.Tensor | None = None,
+) -> torch.Tensor:
+    B, T, Hg, K, V = *q.shape, v.shape[-1]
+    H = v.shape[-2]
+    BT = chunk_size
+
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    o = torch.empty_like(v)
+    if cu_seqlens is None:
+        N, chunk_offsets = B, None
+    else:
+        N = len(cu_seqlens) - 1
+        if chunk_offsets is None:
+            chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), N * H)
+
+    g = g.transpose(1, 2).contiguous()
+    chunk_fwd_kernel_o[grid](
+        q=q,
+        k=k,
+        v=v,
+        h=h,
+        g=g,
+        o=o,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        scale=scale,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=128,
+        BV=128,
+        num_warps=4,
+        num_stages=2,
+    )
+    return o

vllm_ascend/ops/triton/fla/chunk_o_update.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .utils import prepare_chunk_offsets
+
+
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_fwd_kernel_o_update(
+    h,
+    h_update,
+    updated_h_state,
+    cu_seqlens,
+    chunk_offsets,
+    T,
+    H: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H  # splitting by the head of the req
+
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int64)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    # offset calculation
+    updated_h_state += (i_n * H + i_h).to(tl.int64) * K * V
+
+    for i_t in range(NT):
+        i_tg = boh + i_t
+        h_base = h + (i_tg * H + i_h).to(tl.int64) * K * V
+        hupd_base = h_update + ((i_tg + i_n) * H + i_h).to(tl.int64) * K * K
+
+        for i_k in range(tl.cdiv(K, BK)):
+            p_h = tl.make_block_ptr(h_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+            p_hupd = tl.make_block_ptr(hupd_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BK), (1, 0))
+            p_updated_h_state = tl.make_block_ptr(
+                updated_h_state, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)
+            )
+
+            # [BK, BV]
+            b_h = tl.load(p_h, boundary_check=(0, 1))
+            # [BK, BK]
+            b_hupd = tl.load(p_hupd, boundary_check=(0, 1))
+            # [BK, BV]
+            b_updated_h_state = tl.load(p_updated_h_state, boundary_check=(0, 1))
+
+            b_h += tl.dot(b_hupd.to(tl.bfloat16), b_updated_h_state.to(tl.bfloat16))
+            tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_fwd_o_update(
+    q: torch.Tensor,
+    v: torch.Tensor,
+    h: torch.Tensor,
+    h_update: torch.Tensor,
+    updated_h_state: torch.Tensor,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    chunk_offsets: torch.Tensor | None = None,
+) -> torch.Tensor:
+    B, T, Hg, K, V = *q.shape, v.shape[-1]
+    H = v.shape[-2]
+    BT = chunk_size
+
+    if cu_seqlens is None:
+        N, chunk_offsets = B, None
+    else:
+        N = len(cu_seqlens) - 1
+        if chunk_offsets is None:
+            chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), N * H)
+
+    chunk_fwd_kernel_o_update[grid](
+        h=h,
+        h_update=h_update,
+        updated_h_state=updated_h_state,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=128,
+        BV=128,
+        num_warps=4,
+        num_stages=2,
+    )
+    return h

vllm_ascend/ops/triton/fla/chunk_scaled_dot_kkt.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import get_aicore_num
+
+from .utils import prepare_chunk_indices, safe_exp
+
+
+@triton.heuristics(
+    {
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+        "USE_G": lambda args: args["g_cumsum"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T", "B"])
+def chunk_scaled_dot_kkt_fwd_kernel(
+    k,
+    beta,  # [H, B, T]
+    g_cumsum,  # [H, B, T]
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B,
+    H: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_G: tl.constexpr,
+    bh_step: tl.constexpr,
+    task_num: tl.constexpr,
+    num_core: tl.constexpr,
+):
+    bt_stride = B * T
+    core_id = tl.program_id(0)
+
+    for task_id in tl.range(core_id, task_num, num_core):
+        i_t_i = task_id // bh_step
+        i_bh = task_id % bh_step
+        i_b, i_h = i_bh // H, i_bh % H
+        if IS_VARLEN:
+            i_n, i_t = (
+                tl.load(chunk_indices + i_t_i * 2).to(tl.int32),
+                tl.load(chunk_indices + i_t_i * 2 + 1).to(tl.int32),
+            )
+            bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+            T = eos - bos
+        else:
+            bos, eos = i_b * T, i_b * T + T
+            i_t = i_t_i
+        o_t = tl.arange(0, BT)
+        o_t_fp32 = o_t.to(tl.float32)
+
+        p_beta = tl.make_block_ptr(beta + i_h * bt_stride + bos, (T,), (1,), (i_t * BT,), (BT,), (0,))
+        b_beta = tl.load(p_beta, boundary_check=(0,))
+
+        b_A = tl.zeros([BT, BT], dtype=tl.float32)
+        for i_k in range(tl.cdiv(K, BK)):
+            p_k = tl.make_block_ptr(
+                k + (bos * Hg + i_h // (H // Hg)) * K, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+            )
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_A += tl.dot(b_k, tl.trans(b_k))
+
+        if USE_G:
+            p_g = tl.make_block_ptr(g_cumsum + i_h * bt_stride + bos, (T,), (1,), (i_t * BT,), (BT,), (0,))
+            b_g = tl.load(p_g, boundary_check=(0,))
+            b_g_diff = b_g[:, None] - b_g[None, :]
+            b_A *= safe_exp(b_g_diff)
+
+        b_A *= b_beta[:, None]
+        b_A = tl.where(o_t_fp32[:, None] > o_t_fp32[None, :], b_A, 0)
+        p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
+        tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_scaled_dot_kkt_fwd(
+    k: torch.Tensor,
+    beta: torch.Tensor,
+    g_cumsum: torch.Tensor | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+    chunk_size: int = 64,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    r"""
+    Compute beta * K * K^T.
+
+    Args:
+        k (torch.Tensor):
+            The key tensor of shape `[B, T, H, K]`.
+        beta (torch.Tensor):
+            The beta tensor of shape `[B, T, H]`.
+        g (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
+        gk (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
+        cu_seqlens (torch.LongTensor):
+            The cumulative sequence lengths of the input tensor.
+            Default: None
+        chunk_size (int):
+            The chunk size. Default: 64.
+        output_dtype (torch.dtype):
+            The dtype of the output tensor. Default: `torch.float32`
+
+    Returns:
+        beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
+    """
+    B, T, Hg, K = k.shape
+
+    H = beta.shape[-1]
+    BT = chunk_size
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
+
+    num_core = get_aicore_num()
+    bh_step = B * H
+    task_num = NT * bh_step
+
+    from vllm_ascend.device.device_op import DeviceOperator
+
+    A = DeviceOperator.chunk_scaled_dot_kkt_fwd(
+        num_core=num_core,
+        bh_step=bh_step,
+        task_num=task_num,
+        k=k,
+        beta=torch.permute(beta, (2, 0, 1)).contiguous(),
+        g_cumsum=torch.permute(g_cumsum, (2, 0, 1)).contiguous(),
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        B=B,
+        H=H,
+        Hg=Hg,
+        K=K,
+        BT=BT,
+        BK=128,
+    )
+    return A

vllm_ascend/ops/triton/fla/cumsum.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .utils import prepare_chunk_indices
+
+
+@triton.heuristics(
+    {"HAS_SCALE": lambda args: args["scale"] is not None, "IS_VARLEN": lambda args: args["cu_seqlens"] is not None}
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_local_cumsum_scalar_kernel(
+    s,
+    o,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    REVERSE: tl.constexpr,
+    HAS_SCALE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+    CHUNK_SIZE: tl.constexpr = 64,
+):
+    i_block, i_b = tl.program_id(0), tl.program_id(1)
+    N_CHUNKS: tl.constexpr = BLOCK_T // CHUNK_SIZE
+
+    if IS_VARLEN:
+        i_s, i_block = (
+            tl.load(chunk_indices + i_block * 2).to(tl.int32),
+            tl.load(chunk_indices + i_block * 2 + 1).to(tl.int32),
+        )
+        bos, eos = tl.load(cu_seqlens + i_s).to(tl.int32), tl.load(cu_seqlens + i_s + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    if HEAD_FIRST:
+        ptr_s = tl.make_block_ptr(s + bos * H, (H, T), (T, 1), (0, i_block * BLOCK_T), (H, BLOCK_T), (1, 0))
+        ptr_o = tl.make_block_ptr(o + bos * H, (H, T), (T, 1), (0, i_block * BLOCK_T), (H, BLOCK_T), (1, 0))
+        b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32)
+        b_s = tl.reshape(b_s, (H, N_CHUNKS, CHUNK_SIZE))
+        b_s = tl.trans(b_s, (2, 0, 1))
+        b_o = tl.cumsum(b_s, axis=0, reverse=REVERSE)
+        if HAS_SCALE:
+            b_o *= scale
+        b_o = tl.trans(b_o, (2, 0, 1))
+        b_o = tl.reshape(b_o, (H, BLOCK_T))
+    else:
+        ptr_s = tl.make_block_ptr(s + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0))
+        ptr_o = tl.make_block_ptr(o + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0))
+        b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32)
+        b_s = tl.reshape(b_s, (N_CHUNKS, CHUNK_SIZE, H))
+        b_s = tl.trans(b_s, (1, 0, 2))
+        b_o = tl.cumsum(b_s, axis=0, reverse=REVERSE)
+        if HAS_SCALE:
+            b_o *= scale
+        b_o = tl.trans(b_o, (1, 0, 2))
+        b_o = tl.reshape(b_o, (BLOCK_T, H))
+
+    tl.store(ptr_o, b_o.to(s.dtype.element_ty), boundary_check=(0,))
+    return
+
+
+def chunk_local_cumsum_scalar(
+    g,
+    chunk_size,
+    reverse: bool = False,
+    scale: float = None,
+    cu_seqlens: torch.Tensor | None = None,
+    block_indices: torch.Tensor | None = None,
+    head_first: bool = False,
+    output_dtype: torch.Tensor | None = torch.float,
+):
+    if head_first:
+        B, H, T = g.shape
+    else:
+        B, T, H = g.shape
+    assert chunk_size == 2 ** (chunk_size.bit_length() - 1), "chunk_size must be a power of 2"
+    OPTIM_BLOCK_SIZE = triton.next_power_of_2((2**18) // (H * chunk_size))
+    if cu_seqlens is not None and block_indices is None:
+        block_indices = prepare_chunk_indices(cu_seqlens, chunk_size=OPTIM_BLOCK_SIZE)
+    num_blocks = len(block_indices) if cu_seqlens is not None else triton.cdiv(T, OPTIM_BLOCK_SIZE)
+    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    grid = (num_blocks, B)
+    chunk_local_cumsum_scalar_kernel[grid](
+        s=g_org,
+        o=g,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=block_indices,
+        T=T,
+        H=H,
+        BLOCK_T=OPTIM_BLOCK_SIZE,
+        CHUNK_SIZE=chunk_size,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse,
+        num_warps=8,
+        num_stages=3,
+    )
+    return g
+
+
+def chunk_local_cumsum(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    scale: float = None,
+    cu_seqlens: torch.Tensor | None = None,
+    head_first: bool = False,
+    output_dtype: torch.dtype | None = torch.float,
+    **kwargs,
+) -> torch.Tensor:
+    if cu_seqlens is not None:
+        assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided"
+    if len(g.shape) == 3:
+        return chunk_local_cumsum_scalar(
+            g=g,
+            chunk_size=chunk_size,
+            reverse=reverse,
+            scale=scale,
+            cu_seqlens=cu_seqlens,
+            block_indices=kwargs.get("block_indices"),
+            head_first=head_first,
+            output_dtype=output_dtype,
+        )
+    else:
+        raise ValueError(
+            f"chunk_local_cumsum: Unsupported input shape {g.shape}, "
+            f"which should be (B, T, H, D) if `head_first=False` "
+            f"or (B, H, T, D) otherwise"
+        )

vllm_ascend/ops/triton/fla/fused_qkvzba_split_reshape.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# ruff: noqa: E501
+# mypy: ignore-errors
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
+
+MAX_ROWS_PER_ITER = 64
+
+
+@triton.jit(do_not_specialize=["total_rows", "rows_per_vec"])
+def fused_qkvzba_split_reshape_cat_kernel(
+    mixed_qkv,
+    z,
+    b,
+    a,
+    mixed_qkvz,
+    mixed_ba,
+    NUM_HEADS_QK: tl.constexpr,
+    NUM_HEADS_V: tl.constexpr,
+    HEAD_QK: tl.constexpr,
+    HEAD_V: tl.constexpr,
+    total_rows,
+    rows_per_vec,
+    QKVZ_ROW_STRIDE: tl.constexpr,
+    BA_ROW_STRIDE: tl.constexpr,
+    QKV_ROW_STRIDE: tl.constexpr,
+    Z_ROW_STRIDE: tl.constexpr,
+    BA_OUT_ROW_STRIDE: tl.constexpr,
+    ROWS_PER_ITER: tl.constexpr,
+):
+    """
+    Fused kernel to split and reshape mixed QKVZ and BA tensors.
+
+    This kernel performs the following transformations:
+    - Input mixed_qkvz: [num_tokens, num_heads_qk * (Q + K + V + Z)] where each
+      head block contains [Q(HEAD_QK), K(HEAD_QK), V(V_DIM_PER_QK), Z(V_DIM_PER_QK)]
+    - Input mixed_ba: [num_tokens, num_heads_qk * (B + A)] where each head block
+      contains [B(V_HEADS_PER_QK), A(V_HEADS_PER_QK)]
+    - Output mixed_qkv: [num_tokens, Q_all | K_all | V_all] concatenated by type
+    - Output z: [num_tokens, num_heads_v, head_v]
+    - Output b, a: [num_tokens, num_heads_v]
+    """
+    # Each vector core processes a contiguous chunk of rows
+    vec_id = tl.program_id(0)
+
+    V_HEADS_PER_QK: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK
+    V_DIM_PER_QK: tl.constexpr = V_HEADS_PER_QK * HEAD_V
+    QKVZ_DIM_T: tl.constexpr = HEAD_QK * 2 + V_DIM_PER_QK * 2
+    BA_DIM_T: tl.constexpr = V_HEADS_PER_QK * 2
+
+    Q_TOTAL: tl.constexpr = NUM_HEADS_QK * HEAD_QK
+    K_TOTAL: tl.constexpr = NUM_HEADS_QK * HEAD_QK
+
+    row_start = vec_id * rows_per_vec
+    row_end = min(row_start + rows_per_vec, total_rows)
+
+    row_offset = row_start
+
+    iter_count = (row_end - row_start + ROWS_PER_ITER - 1) // ROWS_PER_ITER
+
+    # ========== Main Iteration Loop ==========
+    for _ in tl.range(iter_count):
+        row_indices = tl.arange(0, ROWS_PER_ITER) + row_offset
+        row_mask = row_indices < row_end
+
+        # ========== Head Iteration Loop ==========
+        # Iterate over each Q/K head group to extract and rearrange data
+        for head_id in tl.static_range(NUM_HEADS_QK):
+            # Byte offset to the current head's data block in mixed_qkvz
+            src_head_offset = head_id * QKVZ_DIM_T
+
+            # ----- Q (Query) Extraction -----
+            # Source layout: mixed_qkvz[row, head_id * QKVZ_DIM_T + 0:HEAD_QK]
+            # Dest layout: mixed_qkv[row, head_id * HEAD_QK : (head_id+1) * HEAD_QK]
+            q_range = tl.arange(0, HEAD_QK)
+            q_src = row_indices[:, None] * QKVZ_ROW_STRIDE + src_head_offset + q_range[None, :]
+            q_dst = row_indices[:, None] * QKV_ROW_STRIDE + head_id * HEAD_QK + q_range[None, :]
+            q_data = tl.load(mixed_qkvz + q_src, mask=row_mask[:, None])
+            tl.store(mixed_qkv + q_dst, q_data, mask=row_mask[:, None])
+
+            # ----- K (Key) Extraction -----
+            # Source layout: mixed_qkvz[row, head_id * QKVZ_DIM_T + HEAD_QK : +HEAD_QK]
+            # Dest layout: mixed_qkv[row, Q_TOTAL + head_id * HEAD_QK : ...]
+            # K is stored after Q in the source; in dest, K starts after all Q heads
+            k_src = row_indices[:, None] * QKVZ_ROW_STRIDE + src_head_offset + HEAD_QK + q_range[None, :]
+            k_dst = row_indices[:, None] * QKV_ROW_STRIDE + Q_TOTAL + head_id * HEAD_QK + q_range[None, :]
+            k_data = tl.load(mixed_qkvz + k_src, mask=row_mask[:, None])
+            tl.store(mixed_qkv + k_dst, k_data, mask=row_mask[:, None])
+
+            # ----- V (Value) Extraction -----
+            # Source layout: mixed_qkvz[row, head_id * QKVZ_DIM_T + HEAD_QK*2 : +V_DIM_PER_QK]
+            # Dest layout: mixed_qkv[row, Q_TOTAL + K_TOTAL + head_id * V_DIM_PER_QK : ...]
+            # V follows Q and K in source; in dest, V starts after all Q and K heads
+            v_range = tl.arange(0, V_DIM_PER_QK)
+            v_src = row_indices[:, None] * QKVZ_ROW_STRIDE + src_head_offset + HEAD_QK * 2 + v_range[None, :]
+            v_dst = (
+                row_indices[:, None] * QKV_ROW_STRIDE + Q_TOTAL + K_TOTAL + head_id * V_DIM_PER_QK + v_range[None, :]
+            )
+            v_data = tl.load(mixed_qkvz + v_src, mask=row_mask[:, None])
+            tl.store(mixed_qkv + v_dst, v_data, mask=row_mask[:, None])
+
+            # ----- Z Extraction -----
+            # Source layout: mixed_qkvz[row, head_id * QKVZ_DIM_T + HEAD_QK*2 + V_DIM_PER_QK : ...]
+            # Dest layout: z[row, head_id * V_DIM_PER_QK : (head_id+1) * V_DIM_PER_QK]
+            # Z follows V in source; output z is reshaped to [batch, num_heads_v, head_v]
+            z_src = (
+                row_indices[:, None] * QKVZ_ROW_STRIDE + src_head_offset + HEAD_QK * 2 + V_DIM_PER_QK + v_range[None, :]
+            )
+            z_dst = row_indices[:, None] * Z_ROW_STRIDE + head_id * V_DIM_PER_QK + v_range[None, :]
+            z_data = tl.load(mixed_qkvz + z_src, mask=row_mask[:, None])
+            tl.store(z + z_dst, z_data, mask=row_mask[:, None])
+
+            # ----- B Extraction -----
+            # Source layout: mixed_ba[row, head_id * BA_DIM_T : +V_HEADS_PER_QK]
+            # Dest layout: b[row, head_id * V_HEADS_PER_QK : (head_id+1) * V_HEADS_PER_QK]
+            b_range = tl.arange(0, V_HEADS_PER_QK)
+            ba_head_offset = head_id * BA_DIM_T
+            b_src = row_indices[:, None] * BA_ROW_STRIDE + ba_head_offset + b_range[None, :]
+            b_dst = row_indices[:, None] * BA_OUT_ROW_STRIDE + head_id * V_HEADS_PER_QK + b_range[None, :]
+            b_data = tl.load(mixed_ba + b_src, mask=row_mask[:, None])
+            tl.store(b + b_dst, b_data, mask=row_mask[:, None])
+
+            # ----- A Extraction -----
+            # Source layout: mixed_ba[row, head_id * BA_DIM_T + V_HEADS_PER_QK : ...]
+            # Dest layout: a[row, head_id * V_HEADS_PER_QK : ...] (same as b_dst)
+            # A follows B in source; output layout is same as B
+            a_src = row_indices[:, None] * BA_ROW_STRIDE + ba_head_offset + V_HEADS_PER_QK + b_range[None, :]
+            a_data = tl.load(mixed_ba + a_src, mask=row_mask[:, None])
+            tl.store(a + b_dst, a_data, mask=row_mask[:, None])
+
+        row_offset += ROWS_PER_ITER
+
+
+def fused_qkvzba_split_reshape_cat(
+    mixed_qkvz,
+    mixed_ba,
+    num_heads_qk,
+    num_heads_v,
+    head_qk,
+    head_v,
+):
+    batch, seq_len = mixed_qkvz.shape[0], 1
+    total_rows = batch * seq_len
+
+    v_heads_per_qk = num_heads_v // num_heads_qk
+    v_dim_per_qk = v_heads_per_qk * head_v
+    qkvz_dim_t = head_qk * 2 + v_dim_per_qk * 2
+    ba_dim_t = v_heads_per_qk * 2
+
+    # row stride
+    qkvz_row_stride = num_heads_qk * qkvz_dim_t
+    ba_row_stride = num_heads_qk * ba_dim_t
+    qkv_row_stride = num_heads_qk * head_qk * 2 + num_heads_v * head_v
+    z_row_stride = num_heads_v * head_v
+    ba_out_row_stride = num_heads_v
+
+    qkv_dim_t = num_heads_qk * head_qk * 2 + num_heads_v * head_v
+    mixed_qkv = torch.empty(
+        [batch * seq_len, qkv_dim_t],
+        dtype=mixed_qkvz.dtype,
+        device=mixed_qkvz.device,
+    )
+    z = torch.empty(
+        [batch * seq_len, num_heads_v, head_v],
+        dtype=mixed_qkvz.dtype,
+        device=mixed_qkvz.device,
+    )
+    b = torch.empty(
+        [batch * seq_len, num_heads_v],
+        dtype=mixed_ba.dtype,
+        device=mixed_ba.device,
+    )
+    a = torch.empty(
+        [batch * seq_len, num_heads_v],
+        dtype=mixed_ba.dtype,
+        device=mixed_ba.device,
+    )
+
+    num_vectorcore = get_vectorcore_num()
+
+    grid_size = min(num_vectorcore, total_rows)
+    grid_size = max(1, grid_size)
+
+    rows_per_vec = triton.cdiv(total_rows, grid_size)
+
+    ub_size = 85 * 1024 // mixed_qkvz.element_size()
+
+    elements_per_row = qkvz_row_stride + ba_row_stride + qkv_row_stride + z_row_stride + ba_out_row_stride * 2
+
+    rows_per_iter = max(1, ub_size // elements_per_row)
+    rows_per_iter = triton.next_power_of_2(rows_per_iter)
+    rows_per_iter = min(rows_per_iter, rows_per_vec, MAX_ROWS_PER_ITER)
+
+    grid = (grid_size, 1)
+    fused_qkvzba_split_reshape_cat_kernel[grid](
+        mixed_qkv,
+        z,
+        b,
+        a,
+        mixed_qkvz,
+        mixed_ba,
+        num_heads_qk,
+        num_heads_v,
+        head_qk,
+        head_v,
+        total_rows,
+        rows_per_vec,
+        qkvz_row_stride,
+        ba_row_stride,
+        qkv_row_stride,
+        z_row_stride,
+        ba_out_row_stride,
+        rows_per_iter,
+    )
+    return mixed_qkv, z, b, a

vllm_ascend/ops/triton/fla/l2norm.py

@@ -0,0 +1,66 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fla/ops/l2norm.py
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
+
+
+@triton.jit(do_not_specialize=["eps", "M", "NUM_CHUNKS"])
+def l2norm_fwd_kernel2_loop(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr, NUM_CHUNKS):
+    base_row = tl.program_id(0) * (NUM_CHUNKS * MBLOCK)
+    rindex = tl.arange(0, N)[None, :]
+
+    for chunk in range(NUM_CHUNKS):
+        row_idx = base_row + chunk * MBLOCK + tl.arange(0, MBLOCK)[:, None]
+        xmask = row_idx < M
+
+        xs = tl.load(X + (rindex + N * row_idx), mask=xmask, other=0.0).to(tl.float32)
+        square = xs * xs
+        square_sum = tl.sum(square, 1)[:, None]
+        rsqrt = tl.rsqrt(square_sum + eps)
+
+        tl.store(Y + (rindex + N * row_idx), xs * rsqrt, xmask)
+
+
+def l2norm_fwd(x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None):
+    x_shape_og = x.shape
+    x = x.reshape(-1, x.shape[-1])
+    # allocate output
+    if output_dtype is None:
+        y = torch.empty_like(x)
+    else:
+        y = torch.empty_like(x, dtype=output_dtype)
+    assert y.stride(-1) == 1
+    T, D = x.shape[0], x.shape[-1]
+    # Less than 64KB per feature: enqueue fused kernel
+    MAX_FUSED_SIZE = 65536 // x.element_size()
+    BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
+    if D > BD:
+        raise RuntimeError(f"l2norm_fwd: This layer doesn't support feature dim >= 64KB, got {D}.")
+
+    MBLOCK = 69
+    # M, N = x.shape
+    num_core = get_vectorcore_num()
+    main_bs = triton.cdiv(T, num_core)
+    num_sub_blocks = triton.cdiv(main_bs, MBLOCK)
+    grid = (num_core,)
+    l2norm_fwd_kernel2_loop[grid](
+        X=x,
+        Y=y,
+        eps=eps,
+        M=T,
+        N=D,
+        MBLOCK=MBLOCK,
+        NUM_CHUNKS=num_sub_blocks,
+    )
+
+    return y.view(x_shape_og)

vllm_ascend/ops/triton/fla/layernorm_guard.py

@@ -0,0 +1,198 @@
+# Adapt from https://github.com/fla-org/flash-linear-attention/blob/main/fla/modules/layernorm_gated.py
+# Copyright (c) 2024, Tri Dao.
+# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
+# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate.
+# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling.
+# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+MAX_CORES = 65535
+
+
+@triton.heuristics(
+    {
+        "HAS_BIAS": lambda args: args["B"] is not None,
+        "HAS_Z": lambda args: args["Z"] is not None,
+    }
+)
+@triton.jit
+def layer_norm_fwd_kernel(
+    X,  # pointer to the input
+    Y,  # pointer to the output
+    W,  # pointer to the weights
+    B,  # pointer to the biases
+    Z,  # pointer to the other branch
+    Mean,  # pointer to the mean
+    Rstd,  # pointer to the 1/std
+    stride_x_row,  # how much to increase the pointer when moving by 1 row
+    stride_y_row,
+    stride_z_row,
+    M,  # number of rows in X_base
+    N,  # number of columns in X_base
+    eps,  # epsilon to avoid division by zero
+    BLOCK_N: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    HAS_Z: tl.constexpr,
+    NORM_BEFORE_GATE: tl.constexpr,
+    IS_RMS_NORM: tl.constexpr,
+    N_CORES: tl.constexpr,
+):
+    # Map the program id to the row of X_base and Y_base it should compute.
+    row = tl.program_id(0)
+    group = tl.program_id(1)
+
+    BLOCK_ROWS = M if M < N_CORES else N_CORES
+    n_iters = M // BLOCK_ROWS
+    remain = M % BLOCK_ROWS
+    if row < remain:
+        n_iters = n_iters + 1
+
+    for i in tl.range(n_iters):
+        X_base = X + (i * BLOCK_ROWS * stride_x_row) + row * stride_x_row + group * N
+        Y_base = Y + (i * BLOCK_ROWS * stride_y_row) + row * stride_y_row + group * N
+        if HAS_Z:
+            Z_base = Z + (i * BLOCK_ROWS * stride_z_row) + row * stride_z_row + group * N
+        if not IS_RMS_NORM:
+            Mean_base = Mean + (i * BLOCK_ROWS) + group * M
+        Rstd_base = Rstd + (i * BLOCK_ROWS) + group * M
+        W_base = W + group * N
+        if HAS_BIAS:
+            B_base = B + group * N
+        # Compute mean and variance
+        cols = tl.arange(0, BLOCK_N)
+        x = tl.load(X_base + cols, mask=cols < N, other=0.0).to(tl.float32)
+        if HAS_Z and not NORM_BEFORE_GATE:
+            z = tl.load(Z_base + cols, mask=cols < N).to(tl.float32)
+            x *= z * tl.sigmoid(z)
+        if not IS_RMS_NORM:
+            mean = tl.sum(x, axis=0) / N
+            tl.store(Mean_base + row, mean)
+            xbar = tl.where(cols < N, x - mean, 0.0)
+            var = tl.sum(xbar * xbar, axis=0) / N
+        else:
+            xbar = tl.where(cols < N, x, 0.0)
+            var = tl.sum(xbar * xbar, axis=0) / N
+        rstd = 1 / tl.sqrt(var + eps)
+        tl.store(Rstd_base + row, rstd)
+        # Normalize and apply linear transformation
+        mask = cols < N
+        w = tl.load(W_base + cols, mask=mask).to(tl.float32)
+        if HAS_BIAS:
+            b = tl.load(B_base + cols, mask=mask).to(tl.float32)
+        x_hat = (x - mean) * rstd if not IS_RMS_NORM else x * rstd
+        y = x_hat * w + b if HAS_BIAS else x_hat * w
+        if HAS_Z and NORM_BEFORE_GATE:
+            z = tl.load(Z_base + cols, mask=mask).to(tl.float32)
+            y *= z * tl.sigmoid(z)
+        # Write output
+        tl.store(Y_base + cols, y, mask=mask)
+
+
+def _layer_norm_fwd(
+    x,
+    weight,
+    bias,
+    eps,
+    z=None,
+    out=None,
+    group_size=None,
+    norm_before_gate=True,
+    is_rms_norm=False,
+):
+    M, N = x.shape
+    if group_size is None:
+        group_size = N
+    assert N % group_size == 0
+    ngroups = N // group_size
+    assert x.stride(-1) == 1
+    if z is not None:
+        assert z.stride(-1) == 1
+        assert z.shape == (M, N)
+    assert weight.shape == (N,)
+    assert weight.stride(-1) == 1
+    if bias is not None:
+        assert bias.stride(-1) == 1
+        assert bias.shape == (N,)
+    # allocate output
+    if out is not None:
+        assert out.shape == x.shape
+    else:
+        out = torch.empty_like(x)
+    assert out.stride(-1) == 1
+    mean = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device) if not is_rms_norm else None
+    rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)
+    # Less than 64KB per feature: enqueue fused kernel
+    MAX_FUSED_SIZE = 65536 // x.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size))
+    if group_size > BLOCK_N:
+        raise RuntimeError(f"_layer_norm_fwd: This layer norm doesn't support feature dim >= 64KB, got {group_size}.")
+    # heuristics for number of warps
+    num_warps = min(max(BLOCK_N // 256, 1), 8)
+    grid = (M if M < MAX_CORES else MAX_CORES, ngroups)
+    with torch.npu.device(x.device.index):
+        layer_norm_fwd_kernel[grid](
+            x,
+            out,
+            weight,
+            bias,
+            z,
+            mean,
+            rstd,
+            x.stride(0),
+            out.stride(0),
+            z.stride(0) if z is not None else 0,
+            M,
+            group_size,
+            eps,
+            BLOCK_N=BLOCK_N,
+            NORM_BEFORE_GATE=norm_before_gate,
+            IS_RMS_NORM=is_rms_norm,
+            N_CORES=MAX_CORES,
+            num_warps=num_warps,
+        )
+    return out, mean, rstd
+
+
+class LayerNormFn(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x,
+        weight,
+        bias,
+        z=None,
+        eps=1e-6,
+        group_size=None,
+        norm_before_gate=True,
+        is_rms_norm=False,
+        activation: str = "swish",
+    ):
+        """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z))"""
+
+        x_shape_og = x.shape
+        # reshape input data into 2D tensor
+        x = x.reshape(-1, x.shape[-1])
+        if x.stride(-1) != 1:
+            x = x.contiguous()
+        if z is not None:
+            assert z.shape == x_shape_og
+            z = z.reshape(-1, z.shape[-1])
+            if z.stride(-1) != 1:
+                z = z.contiguous()
+        weight = weight.contiguous()
+        if bias is not None:
+            bias = bias.contiguous()
+        y, mean, rstd = _layer_norm_fwd(
+            x,
+            weight,
+            bias,
+            eps,
+            z=z,
+            group_size=group_size,
+            norm_before_gate=norm_before_gate,
+            is_rms_norm=is_rms_norm,
+        )
+        return y.reshape(x_shape_og)

vllm_ascend/ops/triton/fla/sigmoid_gating.py

@@ -0,0 +1,388 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import os
+
+import torch
+from vllm.triton_utils import tl, tldevice, triton
+
+if os.environ.get("FLA_USE_FAST_OPS", "0") == "1":
+    div = tldevice.fast_dividef
+    exp = tldevice.fast_expf
+    log = tldevice.fast_logf
+    log2 = tldevice.fast_log2f
+else:
+
+    @triton.jit
+    def div_normal(x, y):
+        return x / y
+
+    div = div_normal
+    exp = tl.exp
+    log = tl.log
+    log2 = tl.log2
+
+
+@triton.heuristics(
+    {
+        "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+        "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None,
+        "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["scale", "N", "T", "B"])
+def fused_recurrent_gated_delta_rule_fwd_kernel(
+    q,
+    k,
+    v,
+    g,
+    beta,
+    o,
+    h0,
+    ht,
+    cu_seqlens,
+    ssm_state_indices,
+    num_accepted_tokens,
+    scale,
+    N,  # num of sequences
+    T,  # num of tokens
+    B,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    stride_init_state_token: tl.constexpr,
+    stride_final_state_token: tl.constexpr,
+    stride_indices_seq: tl.constexpr,
+    stride_indices_tok: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,  # whether to use initial state
+    INPLACE_FINAL_STATE: tl.constexpr,  # whether to store final state inplace
+    IS_BETA_HEADWISE: tl.constexpr,  # whether beta is headwise vector or scalar,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    IS_CONTINUOUS_BATCHING: tl.constexpr,
+    IS_SPEC_DECODING: tl.constexpr,
+    IS_KDA: tl.constexpr,
+):
+    i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_n, i_hv = i_nh // HV, i_nh % HV
+    i_h = i_hv // (HV // H)
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
+        all = T
+        T = eos - bos
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        all = B * T
+
+    if T == 0:
+        # no tokens to process for this sequence
+        return
+
+    o_k = i_k * BK + tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+
+    mask_k = o_k < K
+    mask_v = o_v < V
+    mask_h = mask_k[:, None] & mask_v[None, :]
+
+    b_h = tl.zeros([BK, BV], dtype=tl.float32)
+    if USE_INITIAL_STATE:
+        if IS_CONTINUOUS_BATCHING:
+            if IS_SPEC_DECODING:
+                i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
+            else:
+                i_t = 0
+            p_h0 = (
+                h0 + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(tl.int64) * stride_init_state_token
+            )
+        else:
+            p_h0 = h0 + bos * HV * K * V
+        p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
+        b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
+
+    for i_t in range(0, T):
+        p_q = q + (bos * H + i_h) * K + o_k + H * K * i_t
+        p_k = k + (bos * H + i_h) * K + o_k + H * K * i_t
+        p_v = v + (bos * HV + i_hv) * V + o_v + HV * V * i_t
+
+        if IS_BETA_HEADWISE:
+            p_beta = beta + (bos * HV + i_hv) * V + o_v + HV * V * i_t
+        else:
+            p_beta = beta + bos * HV + i_hv + HV * i_t
+
+        if not IS_KDA:
+            p_g = g + bos * HV + i_hv + HV * i_t
+        else:
+            p_gk = g + (bos * HV + i_hv + HV * i_t) * K + o_k
+
+        p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + HV * V * i_t
+
+        b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
+        b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
+        b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
+
+        if USE_QK_L2NORM_IN_KERNEL:
+            b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
+            b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
+        b_q = b_q * scale
+        # [BK, BV]
+        # b_h *= tl.exp(b_g)
+        if not IS_KDA:
+            b_g = tl.load(p_g).to(tl.float32)
+            b_h *= exp(b_g)
+        else:
+            b_gk = tl.load(p_gk).to(tl.float32)
+            b_h *= exp(b_gk[:, None])
+        # [BV]
+        b_v -= tl.sum(b_h * b_k[:, None], 0)
+        if IS_BETA_HEADWISE:
+            b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
+        else:
+            b_beta = tl.load(p_beta).to(tl.float32)
+        b_v *= b_beta
+        # [BK, BV]
+        b_h += b_k[:, None] * b_v[None, :]
+        # [BV]
+        b_o = tl.sum(b_h * b_q[:, None], 0)
+        tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+
+        # keep the states for multi-query tokens
+        if INPLACE_FINAL_STATE:
+            p_ht = (
+                ht + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(tl.int64) * stride_final_state_token
+            )
+        else:
+            p_ht = ht + (bos + i_t) * stride_final_state_token
+        p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
+        tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
+
+
+@triton.heuristics(
+    {
+        "USE_INITIAL_STATE": lambda args: args["h0_source"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def fused_sigmoid_gating_delta_rule_update_kernel(
+    A_log,
+    a,
+    dt_bias,
+    softplus_beta,
+    softplus_threshold,
+    q,
+    k,
+    v,
+    b,
+    o,
+    h0_source,
+    h0_indices,
+    cu_seqlens,
+    scale,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    HV: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    """
+    Fused kernel that combines sigmoid gating computation with recurrent delta rule update.
+    """
+    i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_n, i_hv = i_nh // HV, i_nh % HV
+    i_h = i_hv // (HV // H)
+
+    if IS_VARLEN:
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int64),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int64),
+        )
+        all = T
+        T = eos - bos
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        all = B * T
+
+    o_k = i_k * BK + tl.arange(0, BK)
+    o_v = i_v * BV + tl.arange(0, BV)
+
+    p_q = q + (bos * H + i_h) * K + o_k
+    p_k = k + (bos * H + i_h) * K + o_k
+    p_v = v + (bos * HV + i_hv) * V + o_v
+    p_b = b + bos * HV + i_hv
+    p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
+
+    # Gating computation pointers
+    p_A_log = A_log + i_hv
+    p_a = a + bos * HV + i_hv
+    p_dt_bias = dt_bias + i_hv
+
+    mask_k = o_k < K
+    mask_v = o_v < V
+    mask_h = mask_k[:, None] & mask_v[None, :]
+
+    b_h = tl.zeros([BK, BV], dtype=tl.float32)
+    if USE_INITIAL_STATE:
+        idx = tl.load(h0_indices + i_n)
+        if idx >= 0:
+            p_h0 = h0_source + idx * HV * K * V + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
+            b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
+
+    for i in range(0, T):
+        # Load inputs
+        b_q = tl.load(p_q + i * H * K, mask=mask_k, other=0).to(tl.float32)
+        b_k = tl.load(p_k + i * H * K, mask=mask_k, other=0).to(tl.float32)
+        b_v = tl.load(p_v + i * HV * V, mask=mask_v, other=0).to(tl.float32)
+        b_b = tl.load(p_b + i * HV).to(tl.float32)
+
+        # Compute sigmoid gating
+        # Load gating parameters
+        b_A_log = tl.load(p_A_log).to(tl.float32)
+        b_a = tl.load(p_a + i * HV).to(tl.float32)
+        b_dt_bias = tl.load(p_dt_bias).to(tl.float32)
+
+        # Compute g = -exp(A_log) * softplus(a + dt_bias)
+        x = b_a + b_dt_bias
+        beta_x = softplus_beta * x
+        # Apply softplus with numerical stability
+        softplus_x = tl.where(
+            beta_x <= softplus_threshold,
+            (1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)),
+            x,
+        )
+        b_g = -tl.exp(b_A_log) * softplus_x
+
+        # Compute beta = sigmoid(b)
+        b_beta = 1.0 / (1.0 + tl.exp(-b_b))
+
+        # Apply L2 normalization if enabled
+        if USE_QK_L2NORM_IN_KERNEL:
+            b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6)
+            b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6)
+
+        b_q = b_q * scale
+
+        # Apply gating to hidden state: h *= exp(g)
+        b_h *= tl.exp(b_g)
+
+        # Delta rule: v -= sum(h * k, dim=0)
+        b_v -= tl.sum(b_h * b_k[:, None], 0)
+
+        # Apply beta gating: v *= beta
+        b_v *= b_beta
+
+        # Update hidden state: h += k[:, None] * v[None, :]
+        b_h += b_k[:, None] * b_v[None, :]
+
+        # Compute output: o = sum(h * q, dim=0)
+        b_o = tl.sum(b_h * b_q[:, None], 0)
+        tl.store(p_o + i * HV * V, b_o.to(p_o.dtype.element_ty), mask=mask_v)
+
+        # # Update pointers for next timestep
+        # p_q += H * K
+        # p_k += H * K
+        # p_o += HV * V
+        # p_v += HV * V
+        # p_b += HV
+        # p_a += HV
+
+    # Store final state back to h0_source with bounds checking
+    if USE_INITIAL_STATE:
+        idx = tl.load(h0_indices + i_n)
+        if idx >= 0:
+            p_h0 = h0_source + idx * HV * K * V + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
+            tl.store(p_h0, b_h.to(p_h0.dtype.element_ty), mask=mask_h)
+
+
+def fused_sigmoid_gating_delta_rule_update(
+    A_log: torch.Tensor,
+    a: torch.Tensor,
+    dt_bias: torch.Tensor,
+    softplus_beta: float,
+    softplus_threshold: float,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    b: torch.Tensor,
+    initial_state_source: torch.Tensor,
+    initial_state_indices: torch.Tensor,
+    scale: float = None,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: torch.Tensor = None,
+):
+    """
+    Fused triton implementation of sigmoid gating delta rule update.
+    This function uses a single fused kernel that combines both sigmoid gating computation
+    and the recurrent delta rule update for better performance.
+    """
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    HV = v.shape[2]
+    N = B if cu_seqlens is None else len(cu_seqlens) - 1
+    BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 64)
+    NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
+    assert NK == 1, "NK > 1 is not supported yet"
+    num_stages = 3
+    num_warps = 1
+
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+    else:
+        assert scale > 0, "scale must be positive"
+
+    o = q.new_empty(NK, *v.shape)
+    grid = (NK, NV, N * HV)
+
+    if not initial_state_indices.is_contiguous():
+        initial_state_indices = initial_state_indices.contiguous()
+    if not initial_state_source.is_contiguous():
+        initial_state_source = initial_state_source.contiguous()
+    if not cu_seqlens.is_contiguous():
+        cu_seqlens = cu_seqlens.contiguous()
+
+    fused_sigmoid_gating_delta_rule_update_kernel[grid](
+        A_log=A_log,
+        a=a,
+        dt_bias=dt_bias,
+        softplus_beta=softplus_beta,
+        softplus_threshold=softplus_threshold,
+        q=q,
+        k=k,
+        v=v,
+        b=b,
+        o=o,
+        h0_source=initial_state_source,
+        h0_indices=initial_state_indices,
+        cu_seqlens=cu_seqlens,
+        scale=scale,
+        T=T,
+        B=B,
+        H=H,
+        HV=HV,
+        K=K,
+        V=V,
+        BK=BK,
+        BV=BV,
+        USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    o = o.squeeze(0)
+    return o

vllm_ascend/ops/triton/fla/solve_tril.py

@@ -0,0 +1,404 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from vllm_ascend.ops.triton.triton_utils import extract_slice, insert_slice
+
+from .utils import prepare_chunk_indices
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T", "H"])
+def solve_tril_16x16_kernel(
+    A,
+    Ad,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H,
+    BT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    LARGE_BLOCK_T: tl.constexpr,
+    EXTRACT_SLICE_STRIDE_1: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    A = A + (bos * H + i_h) * BT
+    Ad = Ad + (bos * H + i_h) * 16
+
+    base_t = i_t * LARGE_BLOCK_T
+
+    NTASKS: tl.constexpr = 2
+    N_BLOCKS: tl.constexpr = LARGE_BLOCK_T // 16 // NTASKS
+
+    for taskid in range(0, NTASKS):
+        base_t += taskid * (LARGE_BLOCK_T // NTASKS)
+
+        # use make_block_ptr to reduce vector computation
+        b_A = tl.zeros((N_BLOCKS, 16, 16), dtype=tl.float32)
+        for blkid in range(0, N_BLOCKS):
+            row_start_o = base_t + blkid * 16
+            col_start_o = row_start_o % BT
+
+            # 1 Create in-block offset
+            offs_rows_in_block = tl.arange(0, 16)
+            offs_cols_in_block = tl.arange(0, 16)
+
+            # 2 Calculate the pointer of each element
+            ptr_A_subrec16 = (
+                A
+                + row_start_o * H * BT
+                + col_start_o
+                + offs_rows_in_block[:, None] * H * BT
+                + offs_cols_in_block[None, :]
+            )
+
+            # 3 Create a mask to prevent out-of-bounds access
+            global_rows = row_start_o + offs_rows_in_block[:, None]
+            global_cols = col_start_o + offs_cols_in_block[None, :]
+            load_mask = (global_rows < T) & (global_cols < BT)
+
+            # 4 Use mask to safely load data
+            b_A_subrec16 = tl.load(ptr_A_subrec16, mask=load_mask, other=0.0).to(tl.float32)
+            b_A = insert_slice(
+                ful=b_A,
+                sub=b_A_subrec16[None, :, :],  # (1, 16, 16)
+                offsets=[blkid, 0, 0],
+                sizes=[1, 16, 16],
+                strides=[1, 1, 1],
+            )
+
+        local_ori_A = tl.trans(b_A, (1, 0, 2))
+        local_ori_A = tl.reshape(local_ori_A, (16, 16 * N_BLOCKS))
+
+        # Convert mask into matrix multiplication to avoid for loops ub oom
+        tmp = tl.arange(0, 16).to(tl.float32)
+        rows = tmp[:, None]
+        cols = tmp[None, :]
+        is_lower = (rows > cols).to(b_A.dtype)
+        b_A = -b_A * is_lower
+
+        # for loop to update N_BLOCKS row vector
+        for i in range(1, 16):
+            nblks_vec16 = -extract_slice(local_ori_A, (i, 0), (1, 16 * N_BLOCKS), (EXTRACT_SLICE_STRIDE_1, 1))
+            b_a = tl.reshape(nblks_vec16, (N_BLOCKS, 16))
+
+            dot_tmp = tl.trans(b_a[:, :, None] * b_A, (1, 0, 2))
+            dot_product = tl.sum(dot_tmp, 0)
+            b_a = b_a + dot_product
+
+            b_a_new_expanded = b_a[:, None, :]
+            b_A = insert_slice(
+                ful=b_A, sub=b_a_new_expanded, offsets=[0, i, 0], sizes=[N_BLOCKS, 1, 16], strides=[1, 1, 1]
+            )
+
+        on_diagonal = rows == cols
+        b_A = tl.where(on_diagonal, b_A + 1.0, b_A)
+
+        b_A = tl.reshape(b_A, (N_BLOCKS * 16, 16))
+        p_Ai = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (base_t, 0), (N_BLOCKS * 16, 16), (1, 0))
+
+        # 1 Create in-block offset
+        offs_rows_to_store = tl.arange(0, N_BLOCKS * 16)
+        offs_cols_to_store = tl.arange(0, 16)
+
+        # 2 Calculate the pointer of each element
+        p_Ai = Ad + base_t * H * 16 + 0 + offs_rows_to_store[:, None] * H * 16 + offs_cols_to_store[None, :]
+        # 3 Create a mask to prevent out-of-bounds access, only check rows
+        global_store_rows = base_t + offs_rows_to_store[:, None]
+        store_mask = global_store_rows < T
+        # 4 use mask to save data safely
+        tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=store_mask)
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T", "H"])
+def merge_16x16_to_32x32_inverse_kernel(
+    A,
+    Ad,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H,
+    BT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    A += (bos * H + i_h) * 32
+    Ad += (bos * H + i_h) * 16
+    Ai += (bos * H + i_h) * 32
+
+    p_A_21 = tl.make_block_ptr(A, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0))
+    p_Ad_11 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 32, 0), (16, 16), (1, 0))
+    p_Ad_22 = tl.make_block_ptr(Ad, (T, 16), (H * 16, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0))
+    p_Ai_11 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32, 0), (16, 16), (1, 0))
+    p_Ai_22 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 16), (16, 16), (1, 0))
+    p_Ai_21 = tl.make_block_ptr(Ai, (T, 32), (H * 32, 1), (i_t * 32 + 16, 0), (16, 16), (1, 0))
+
+    A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
+    Ai_11 = tl.load(p_Ad_11, boundary_check=(0, 1)).to(tl.float32)
+    Ai_22 = tl.load(p_Ad_22, boundary_check=(0, 1)).to(tl.float32)
+    Ai_21 = -tl.dot(
+        tl.dot(Ai_22, A_21, input_precision="ieee"),
+        Ai_11,
+        input_precision="ieee",
+    )
+    tl.store(
+        p_Ai_11,
+        Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"),
+        boundary_check=(0, 1),
+    )
+    tl.store(
+        p_Ai_22,
+        Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"),
+        boundary_check=(0, 1),
+    )
+    tl.store(
+        p_Ai_21,
+        Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"),
+        boundary_check=(0, 1),
+    )
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T", "H"])
+def merge_16x16_to_64x64_inverse_kernel(
+    A,
+    Ad,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H,
+    BT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t_val = (
+            tl.load(chunk_indices + i_t * 2).to(tl.int32),
+            tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
+        )
+        bos, eos = (
+            tl.load(cu_seqlens + i_n).to(tl.int32),
+            tl.load(cu_seqlens + i_n + 1).to(tl.int32),
+        )
+        T = eos - bos
+        i_t = i_t_val
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    # Base pointers (already offset by batch and head)
+    A += (bos * H + i_h) * 64
+    Ad += (bos * H + i_h) * 16
+    Ai += (bos * H + i_h) * 64
+
+    # load Ai_22 (Ad block at row i_t * 64 + 16, col 0, 16 * 16)
+    offs_m = i_t * 64 + 16 + tl.arange(0, 16)
+    offs_n = tl.arange(0, 16)
+    mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16)
+    ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :]
+    Ai_22 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32)
+
+    # load A_21 (A block at row i_t * 64 + 16, col 0, 16 * 16)
+    mask_A = (offs_m[:, None] < T) & (offs_n[None, :] < 64)
+    ptr_A = A + offs_m[:, None] * (H * 64) + offs_n[None, :]
+    A_21 = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32)
+    tmp = tl.dot(Ai_22, A_21, input_precision="ieee")
+
+    # load Ai_11 (Ad block at row i_t * 64, col 0, 16 * 16)
+    offs_m = i_t * 64 + tl.arange(0, 16)
+    offs_n = tl.arange(0, 16)
+    mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16)
+    ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :]
+    Ai_11 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32)
+
+    Ai_21 = -tl.dot(tmp, Ai_11, input_precision="ieee")
+
+    # load Ai_44 (Ad block at row i_t * 64 + 48, col 0, 16 * 16)
+    offs_m = i_t * 64 + 48 + tl.arange(0, 16)
+    offs_n = tl.arange(0, 16)
+    mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16)
+    ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :]
+    Ai_44 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32)
+
+    # load A_43 (Ad block at row i_t * 64 + 48, col 32, 16 * 16)
+    offs_n = 32 + tl.arange(0, 16)
+    mask_A = (offs_m[:, None] < T) & (offs_n[None, :] < 64)
+    ptr_A = A + offs_m[:, None] * (H * 64) + offs_n[None, :]
+    A_43 = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32)
+    tmp = tl.dot(Ai_44, A_43, input_precision="ieee")
+
+    # load Ai_33 (Ad block at row i_t * 64 + 32, col 0, 16 * 16)
+    offs_m = i_t * 64 + 32 + tl.arange(0, 16)
+    offs_n = tl.arange(0, 16)
+    mask_Ad = (offs_m[:, None] < T) & (offs_n[None, :] < 16)
+    ptr_Ad = Ad + offs_m[:, None] * (H * 16) + offs_n[None, :]
+    Ai_33 = tl.load(ptr_Ad, mask=mask_Ad, other=0.0).to(tl.float32)
+
+    Ai_43 = -tl.dot(tmp, Ai_33, input_precision="ieee")
+
+    # build Ai_22_32 (32 * 32)
+    Ai_22_32 = tl.zeros((32, 32), tl.float32)
+    Ai_22_32 = insert_slice(Ai_22_32, Ai_33, (0, 0), (16, 16), (1, 1))
+    Ai_22_32 = insert_slice(Ai_22_32, Ai_44, (16, 16), (16, 16), (1, 1))
+    Ai_22_32 = insert_slice(Ai_22_32, Ai_43, (16, 0), (16, 16), (1, 1))
+
+    # load A_21_32 (A block at row i_t * 64 + 32, col 0, 32 * 32)
+    offs_m = i_t * 64 + 32 + tl.arange(0, 32)
+    offs_n = tl.arange(0, 32)
+    mask_A = (offs_m[:, None] < T) & (offs_n[None, :] < 64)
+    ptr_A = A + offs_m[:, None] * (H * 64) + offs_n[None, :]
+    A_21_32 = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32)
+    tmp = tl.dot(Ai_22_32, A_21_32, input_precision="ieee")
+
+    # build Ai_11_32 (32 * 32)
+    Ai_11_32 = tl.zeros((32, 32), tl.float32)
+    Ai_11_32 = insert_slice(Ai_11_32, Ai_11, (0, 0), (16, 16), (1, 1))
+    Ai_11_32 = insert_slice(Ai_11_32, Ai_22, (16, 16), (16, 16), (1, 1))
+    Ai_11_32 = insert_slice(Ai_11_32, Ai_21, (16, 0), (16, 16), (1, 1))
+
+    Ai_21_32 = -tl.dot(tmp, Ai_11_32, input_precision="ieee")
+
+    # store Ai_11_32 to (i_t * 64, 0)
+    offs_m = i_t * 64 + tl.arange(0, 32)
+    offs_n = tl.arange(0, 32)
+    mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < 64)
+    ptr_Ai = Ai + offs_m[:, None] * (H * 64) + offs_n[None, :]
+    tl.store(ptr_Ai, Ai_11_32.to(ptr_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=mask_store)
+
+    # store Ai_22_32 to (i_t * 64 + 32, 32)
+    offs_m = i_t * 64 + 32 + tl.arange(0, 32)
+    offs_n = 32 + tl.arange(0, 32)
+    mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < 64)
+    ptr_Ai = Ai + offs_m[:, None] * (H * 64) + offs_n[None, :]
+    tl.store(ptr_Ai, Ai_22_32.to(ptr_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=mask_store)
+
+    # store Ai_21_32 to (i_t * 64 + 32, 32)
+    offs_n = tl.arange(0, 32)
+    mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < 64)
+    ptr_Ai = Ai + offs_m[:, None] * (H * 64) + offs_n[None, :]
+    tl.store(ptr_Ai, Ai_21_32.to(ptr_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=mask_store)
+
+    # zero out the upper-right 32 * 32 block (rows 0 ~ 31, cols 32 ~ 63)
+    offs_m = i_t * 64 + tl.arange(0, 32)
+    offs_n = 32 + tl.arange(0, 32)
+    mask_store = (offs_m[:, None] < T) & (offs_n[None, :] < BT)
+    ptr_Ai = Ai + offs_m[:, None] * (H * BT) + offs_n[None, :]
+    zero_block = tl.zeros((32, 32), dtype=ptr_Ai.dtype.element_ty)
+    tl.store(ptr_Ai, zero_block, mask=mask_store)
+
+
+def solve_tril(
+    A: torch.Tensor,
+    cu_seqlens: torch.Tensor | None = None,
+    chunk_indices_large_block: torch.Tensor | None = None,
+    chunk_indices_bt: torch.Tensor | None = None,
+    output_dtype: torch.dtype = torch.float,
+) -> torch.Tensor:
+    """
+    Compute the inverse of the matrix I + A
+    A should be strictly lower triangular, i.e., A.triu() == 0.
+
+    Args:
+        A (torch.Tensor):
+            [B, T, H, BT], where BT should only be 16, 32, or 64.
+        cu_seqlens (torch.Tensor):
+            The cumulative sequence lengths of the input tensor. Default: `None`.
+        output_dtype (torch.dtype):
+            The dtype of the output tensor. Default: `torch.float`.
+            If `None`, the output dtype will be the same as the input dtype.
+
+    Returns:
+        (I + A)^-1 with the same shape as A
+    """
+    assert A.shape[-1] in [16, 32, 64]
+
+    B, T, H, BT = A.shape
+    Ad = torch.empty(B, T, H, 16, device=A.device, dtype=torch.float if BT != 16 else output_dtype)
+
+    LARGE_BLOCK_T = 608 * 2
+
+    if cu_seqlens is not None and chunk_indices_large_block is None:
+        chunk_indices_large_block = prepare_chunk_indices(cu_seqlens, LARGE_BLOCK_T)
+    chunk_indices = chunk_indices_large_block
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, LARGE_BLOCK_T)
+
+    from vllm_ascend.device.device_op import DeviceOperator
+
+    DeviceOperator.solve_tril_16x16(
+        A=A,
+        Ad=Ad,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        BT=BT,
+        LARGE_BLOCK_T=LARGE_BLOCK_T,
+        NT=NT,
+        B=B,
+    )
+
+    if BT == 16:
+        return Ad
+
+    Ai = torch.empty(B, T, H, BT, device=A.device, dtype=output_dtype)
+    merge_fn = merge_16x16_to_32x32_inverse_kernel if BT == 32 else merge_16x16_to_64x64_inverse_kernel
+    if cu_seqlens is not None and chunk_indices_bt is None:
+        chunk_indices_bt = prepare_chunk_indices(cu_seqlens, BT)
+    chunk_indices = chunk_indices_bt
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+
+    merge_fn[NT, B * H](
+        A=A,
+        Ad=Ad,
+        Ai=Ai,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        BT=BT,
+        num_warps=4,
+        num_stages=3,
+    )
+
+    return Ai

vllm_ascend/ops/triton/fla/utils.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# ruff: noqa: E501
+import contextlib
+import functools
+from collections.abc import Callable
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
+    return cu_seqlens[1:] - cu_seqlens[:-1]
+
+
+def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
+    indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()])
+    return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)
+
+
+def prepare_final_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
+    indices = triton.cdiv(prepare_lens(cu_seqlens), chunk_size) + 1
+    return torch.cumsum(indices, 0) - 1
+
+
+def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
+    return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1)
+
+
+def prepare_update_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
+    return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size) + 1]).cumsum(-1)
+
+
+def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
+    """
+    A decorator to make sure all input tensors are contiguous and set the device based on input tensors.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        contiguous_args = (i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args)
+        contiguous_kwargs = {k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}
+
+        tensor = None
+        for arg in args:
+            if isinstance(arg, torch.Tensor):
+                tensor = arg
+                break
+        if tensor is None:
+            for value in kwargs.values():
+                if isinstance(value, torch.Tensor):
+                    tensor = value
+                    break
+
+        if tensor is not None:
+            ctx = torch.npu.device(tensor.device.index)
+        else:
+            ctx = contextlib.nullcontext()
+
+        with ctx:
+            return fn(*contiguous_args, **contiguous_kwargs)
+
+    return wrapper
+
+
+@triton.jit
+def safe_exp(x):
+    return tl.exp(tl.where(x <= 0, x, float("-inf")))
+
+
+@triton.jit(do_not_specialize=["inner_size", "row_stride"])
+def _clear_ssm_states_kernel(
+    states_ptr,
+    has_initial_state_ptr,
+    inner_size,
+    row_stride,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row_idx = tl.program_id(axis=0)
+    col_block_idx = tl.program_id(axis=1)
+
+    has_state = tl.load(has_initial_state_ptr + row_idx).to(tl.int1)
+    if has_state:
+        return
+
+    cols = col_block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = cols < inner_size
+    row_ptr = states_ptr + row_idx * row_stride + cols
+    tl.store(row_ptr, tl.zeros((BLOCK_SIZE,), dtype=states_ptr.dtype.element_ty), mask=mask)
+
+
+def clear_ssm_states(ssm_states: torch.Tensor, has_initial_state: torch.Tensor) -> None:
+    """Zero out specific rows for the SSM states
+
+    Args:
+        ssm_states (torch.Tensor): input SSM states
+        has_initial_state (torch.Tensor): indicates whether the row has initial states already
+    """
+    if ssm_states.numel() == 0:
+        return
+
+    if has_initial_state.device != ssm_states.device:
+        has_initial_state = has_initial_state.to(ssm_states.device, non_blocking=True)
+    if has_initial_state.dtype != torch.bool:
+        has_initial_state = has_initial_state.to(torch.bool)
+
+    has_initial_state = has_initial_state.reshape(-1).contiguous()
+    num_rows = ssm_states.shape[0]
+    if num_rows == 0:
+        return
+    if has_initial_state.numel() != num_rows:
+        raise ValueError(
+            f"clear_ssm_states: has_initial_state size mismatch: expected {num_rows}, got {has_initial_state.numel()}"
+        )
+    inner_size = ssm_states.numel() // num_rows
+    if inner_size == 0:
+        return
+
+    block_size = 4096
+    grid = (num_rows, triton.cdiv(inner_size, block_size))
+    _clear_ssm_states_kernel[grid](
+        ssm_states,
+        has_initial_state,
+        inner_size,
+        ssm_states.stride(0),
+        BLOCK_SIZE=block_size,
+    )

vllm_ascend/ops/triton/fla/wy_fast.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
+#
+# This file contains code copied from the flash-linear-attention project.
+# The original source code was licensed under the MIT license and included
+# the following copyright notice:
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# ruff: noqa: E501
+# mypy: ignore-errors
+
+import torch
+from vllm.triton_utils import tl, triton
+
+from .utils import prepare_chunk_indices
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T", "H", "Hg", "K", "V"])
+def recompute_w_u_fwd_kernel(
+    k,
+    v,
+    beta,
+    w,
+    u,
+    A,
+    g,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H,
+    Hg,
+    K,
+    V,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    T_max = T
+    i_t_o = tl.program_id(0)
+
+    for i_bh in range(H):
+        i_b, i_h = i_bh // H, i_bh % H
+        if IS_VARLEN:
+            i_n, i_t = (
+                tl.load(chunk_indices + i_t_o * 2).to(tl.int32),
+                tl.load(chunk_indices + i_t_o * 2 + 1).to(tl.int32),
+            )
+            bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+            T = eos - bos
+        else:
+            bos, eos = i_b * T, i_b * T + T
+
+        offs_t = tl.arange(0, BT)
+        global_offs_t = i_t * BT + offs_t
+        mask_t = global_offs_t < T
+
+        offs_t_2d = global_offs_t[:, None]
+        offs_bt = tl.arange(0, BT)[None, :]
+        ptr_A = A + (bos * H + i_h) * BT + offs_t_2d * (H * BT) + offs_bt * 1
+        mask_A = mask_t[:, None]
+        b_A = tl.load(ptr_A, mask=mask_A, other=0.0).to(tl.float32)
+
+        ptr_g = g + bos + i_h * T_max + global_offs_t
+        b_g = tl.exp(tl.load(ptr_g, mask=mask_t, other=0.0)).to(tl.float32)
+
+        ptr_beta = beta + bos + i_h * T_max + global_offs_t
+        b_beta = tl.load(ptr_beta, mask=mask_t, other=0.0).to(tl.float32)
+
+        for i_v in range(tl.cdiv(V, BV)):
+            offs_v = i_v * BV + tl.arange(0, BV)[None, :]
+            mask_v = (mask_t[:, None]) & (offs_v < V)
+
+            ptr_v = v + (bos * H + i_h) * V + offs_t_2d * (H * V) + offs_v * 1
+            b_v = tl.load(ptr_v, mask=mask_v, other=0.0).to(tl.float32)
+
+            b_vb = b_v * b_beta[:, None]
+            b_u = tl.dot(b_A, b_vb, allow_tf32=False)
+
+            ptr_u = u + (bos * H + i_h) * V + offs_t_2d * (H * V) + offs_v * 1
+            tl.store(ptr_u, b_u.to(ptr_u.dtype.element_ty), mask=mask_v)
+
+        for i_k in range(tl.cdiv(K, BK)):
+            offs_k = i_k * BK + tl.arange(0, BK)[None, :]
+            mask_k = (mask_t[:, None]) & (offs_k < K)
+            ptr_k = k + (bos * Hg + i_h // (H // Hg)) * K + offs_t_2d * (Hg * K) + offs_k * 1
+            b_k = tl.load(ptr_k, mask=mask_k, other=0.0).to(tl.float32)
+
+            b_kb = b_k * b_beta[:, None] * b_g[:, None]
+            b_w = tl.dot(b_A, b_kb)
+
+            ptr_w = w + (bos * H + i_h) * K + offs_t_2d * (H * K) + offs_k * 1
+            tl.store(ptr_w, b_w.to(ptr_w.dtype.element_ty), mask=mask_k)
+
+
+def recompute_w_u_fwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    g_cumsum: torch.Tensor,
+    A: torch.Tensor,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_indices: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, Hg, K, V = *k.shape, v.shape[-1]
+    H = v.shape[-2]
+    BT = A.shape[-1]
+
+    if cu_seqlens is not None and chunk_indices is None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    BK = 64
+    BV = 64
+
+    u = torch.empty_like(v)
+    w = k.new_empty(B, T, H, K)
+    beta = beta.transpose(1, 2).contiguous()
+    g_cumsum = g_cumsum.transpose(1, 2).contiguous()
+    recompute_w_u_fwd_kernel[(NT, B)](
+        k=k,
+        v=v,
+        beta=beta,
+        w=w,
+        u=u,
+        A=A,
+        g=g_cumsum,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        num_warps=4,
+        num_stages=3,
+    )
+    return w, u