Ascend refacter fused module. (#136)

Vendor

Improvements

Narrow down the patch scope for the module.

Co-authored-by: cyber-pioneer <116002591+cyber-pioneer@users.noreply.github.com>

Author: Yanguan619

vllm_fl/dispatch/backends/vendor/ascend/impl/attention.py

@@ -548,8 +548,9 @@ def _get_fia_params(
         attn_metadata: AscendMetadata,
     ):
         """Get parameters for fused_infer_attention."""
-        block_size = 128
+
         if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
+            block_size = 128
             block_table = None
             actual_seq_lengths_kv = attn_metadata.actual_seq_lengths_q
         elif attn_metadata.attn_state == AscendAttentionState.PrefillCacheHit:
@@ -560,13 +561,18 @@ def _get_fia_params(
             value = self.value_cache.view(num_block, block_size, -1)
             actual_seq_lengths_kv = attn_metadata.seq_lens_list
         elif attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
+            # num_block, block_size, _, _ = self.key_cache.shape
+            # key = self.key_cache.view(num_block, block_size, -1)
+            # value = self.value_cache.view(num_block, block_size, -1)
             key = self.key_cache.view(-1, block_size, 256)
             value = self.value_cache.view(-1, block_size, 256)
             block_table = attn_metadata.block_tables
             actual_seq_lengths_kv = attn_metadata.seq_lens_list
         else:
             # ChunkedPrefill
             # num_block, block_size, _, _ = self.key_cache.shape
+            # key = self.key_cache.view(num_block, block_size, -1)
+            # value = self.value_cache.view(num_block, block_size, -1)
             key = self.key_cache.view(-1, block_size, 256)
             value = self.value_cache.view(-1, block_size, 256)
             block_table = attn_metadata.block_tables

vllm_fl/dispatch/backends/vendor/ascend/impl/fla/__init__.py

@@ -1 +1,6 @@
 # Copyright (c) 2026 BAAI. All rights reserved.
+from vllm_fl.dispatch.backends.vendor.ascend.impl.fla.chunk import (
+    chunk_gated_delta_rule as chunk_gated_delta_rule_npu,
+)
+
+__all__ = ["chunk_gated_delta_rule_npu"]

vllm_fl/dispatch/backends/vendor/ascend/impl/fla/chunk.py

@@ -13,66 +13,10 @@
 
 import torch
 from einops import rearrange
-from vllm.model_executor.layers.fla.ops.utils import SUPPRESS_LEVEL
 
-from .chunk_delta_h import chunk_gated_delta_rule_fwd_h
-from .chunk_o import chunk_fwd_o
-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
-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: Optional[torch.LongTensor] = None,
-):
-    g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)
-    # 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, output_dtype=torch.float32
-    )
-    A = solve_tril(A=A, cu_seqlens=cu_seqlens, 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,
-    )
-    h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
-        k=k,
-        w=w,
-        u=u,
-        g=g,
-        initial_state=initial_state,
-        output_final_state=output_final_state,
-        cu_seqlens=cu_seqlens,
-    )
-    o = chunk_fwd_o(
-        q=q,
-        k=k,
-        v=v_new,
-        h=h,
-        g=g,
-        scale=scale,
-        cu_seqlens=cu_seqlens,
-    )
-    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
-
+from flag_gems.runtime.backend._ascend.fla import chunk_gated_delta_rule_fwd
 
 class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
     @staticmethod