support flash_attention

Author: Shl1015CS

include/llaisys/ops.h

@@ -13,6 +13,18 @@ __C {
     __export void llaisysROPE(llaisysTensor_t out, llaisysTensor_t in, llaisysTensor_t pos_ids, float theta);
     __export void llaisysSelfAttention(llaisysTensor_t attn_val, llaisysTensor_t q, llaisysTensor_t k, llaisysTensor_t v, float scale);
     __export void llaisysSwiGLU(llaisysTensor_t out, llaisysTensor_t gate, llaisysTensor_t up);
+    
+    // Flash Attention with causal masking and GQA support
+    // Behavior matches torch.nn.functional.scaled_dot_product_attention
+    __export void llaisysFlashAttention(
+        llaisysTensor_t out,           // Output tensor [batch, seq_len, num_heads, head_dim]
+        llaisysTensor_t q,             // Query [batch, seq_len, num_q_heads, head_dim]
+        llaisysTensor_t k,             // Key [batch, kv_len, num_kv_heads, head_dim]
+        llaisysTensor_t v,             // Value [batch, kv_len, num_kv_heads, head_dim]
+        float scale,                   // Scale factor (default: 1/sqrt(head_dim))
+        bool is_causal,               // Enable causal masking
+        bool enable_gqa               // Enable Grouped Query Attention
+    );
 }
 
 #endif

python/llaisys/libllaisys/ops.py

@@ -1,5 +1,5 @@
 from .tensor import llaisysTensor_t
-from ctypes import c_float
+from ctypes import c_float, c_bool
 
 def load_ops(lib):
     lib.llaisysAdd.argtypes = [llaisysTensor_t, llaisysTensor_t, llaisysTensor_t]
@@ -34,3 +34,14 @@ def load_ops(lib):
 
     lib.llaisysSwiGLU.argtypes = [llaisysTensor_t, llaisysTensor_t, llaisysTensor_t]
     lib.llaisysSwiGLU.restype = None
+
+    lib.llaisysFlashAttention.argtypes = [
+        llaisysTensor_t,  # out
+        llaisysTensor_t,  # q
+        llaisysTensor_t,  # k
+        llaisysTensor_t,  # v
+        c_float,          # scale
+        c_bool,           # is_causal
+        c_bool            # enable_gqa
+    ]
+    lib.llaisysFlashAttention.restype = None

python/llaisys/models/qwen2.py

@@ -17,12 +17,19 @@ class Qwen2:
 
     DEFAULT_MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 
-    def __init__(self, model_path: Optional[Union[str, Path]] = None, device: DeviceType = DeviceType.CPU):
+    def __init__(
+        self, 
+        model_path: Optional[Union[str, Path]] = None, 
+        device: DeviceType = DeviceType.CPU,
+        max_seq_len: Optional[int] = None
+    ):
         """Initialize Qwen2 model.
         
         Args:
             model_path: Path to model directory. If None, downloads default model.
             device: Device type for inference.
+            max_seq_len: Maximum sequence length. If None, uses model's default.
+                        Set this to a larger value (e.g., 8192, 16384) for longer contexts.
             
         Raises:
             ValueError: If unsupported device is specified.
@@ -40,6 +47,13 @@ def __init__(self, model_path: Optional[Union[str, Path]] = None, device: Device
         config = self._load_config()
         self._init_model_params(config)
 
+        # Override max sequence length if specified
+        if max_seq_len is not None:
+            if max_seq_len <= 0:
+                raise ValueError(f"max_seq_len must be positive, got {max_seq_len}")
+            print(f"[Qwen2] Overriding max_position_embeddings: {self.max_position_embeddings} → {max_seq_len}")
+            self.max_position_embeddings = max_seq_len
+        
         self.data_type = DataType.F32 if device == DeviceType.CPU else DataType.BF16
 
         self._create_model()

python/llaisys/ops.py

@@ -1,6 +1,6 @@
 from .libllaisys import LIB_LLAISYS
 from .tensor import Tensor
-from ctypes import c_float, c_int
+from ctypes import c_float, c_int, c_bool
 
 
 class Ops:
@@ -53,3 +53,36 @@ def self_attention(attn_val: Tensor, q: Tensor, k: Tensor, v: Tensor, scale: flo
     @staticmethod
     def swiglu(out: Tensor, gate: Tensor, up: Tensor):
         LIB_LLAISYS.llaisysSwiGLU(out.lib_tensor(), gate.lib_tensor(), up.lib_tensor())
+
+    @staticmethod
+    def flash_attention(
+        out: Tensor, 
+        q: Tensor, 
+        k: Tensor, 
+        v: Tensor, 
+        scale: float = 0.0,
+        is_causal: bool = False,
+        enable_gqa: bool = False
+    ):
+        """
+        Flash Attention with causal masking and GQA support.
+        Behavior matches torch.nn.functional.scaled_dot_product_attention.
+        
+        Args:
+            out: Output tensor [batch, seq_len, num_q_heads, head_dim]
+            q: Query tensor [batch, seq_len, num_q_heads, head_dim]
+            k: Key tensor [batch, kv_len, num_kv_heads, head_dim]
+            v: Value tensor [batch, kv_len, num_kv_heads, head_dim]
+            scale: Scale factor (default: 1/sqrt(head_dim))
+            is_causal: Enable causal masking
+            enable_gqa: Enable Grouped Query Attention
+        """
+        LIB_LLAISYS.llaisysFlashAttention(
+            out.lib_tensor(),
+            q.lib_tensor(),
+            k.lib_tensor(),
+            v.lib_tensor(),
+            c_float(scale),
+            c_bool(is_causal),
+            c_bool(enable_gqa)
+        )