zyra-network 0.0.1__tar.gz

zyra_network-0.0.1/PKG-INFO

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: zyra-network
+Version: 0.0.1
+Summary: A local agentic AI network that earns Proof of Useful Work (PoUW).
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: PySide6>=6.5.0
+Requires-Dist: PyYAML>=6.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: psutil>=5.9.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: httpx
+Requires-Dist: openai
+Requires-Dist: pyautogui
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pyinstaller; extra == "dev"
+
+# MY-AI
+
+A locally built, trained, and executed Artificial Intelligence desktop application for Windows.
+
+## Tujuan Project
+Proyek ini dibuat untuk membangun neural network Transformer dan seluruh sistem pendukung AI dari awal tanpa bergantung pada layanan cloud, API eksternal, atau model pre-trained yang sudah ada (seperti OpenAI, Claude, Llama, dll). 
+Tujuan akhirnya adalah memiliki AI Assistant pribadi yang berjalan sepenuhnya lokal di komputer pengguna.
+
+## Requirements
+- Windows 64-bit
+- Python 3.10+
+- Hardware minimal: CPU + RAM (Rekomendasi: NVIDIA GPU dengan CUDA support)
+
+## Setup Virtual Environment
+```powershell
+python -m venv venv
+.\venv\Scripts\activate
+```
+
+## Install Dependencies
+Project ini menggunakan `pyproject.toml` untuk manajemen dependensi minimal.
+```powershell
+pip install -e .[dev]
+```
+Dependensi utama:
+- PySide6 (Desktop GUI)
+- PyYAML (Configuration)
+- PyTorch (Hardware detection, Tensor/Autograd)
+
+## Menjalankan Aplikasi
+```powershell
+python run.py
+```
+
+## Menjalankan Test
+```powershell
+pytest tests/
+```

zyra_network-0.0.1/README.md

@@ -0,0 +1,38 @@
+# MY-AI
+
+A locally built, trained, and executed Artificial Intelligence desktop application for Windows.
+
+## Tujuan Project
+Proyek ini dibuat untuk membangun neural network Transformer dan seluruh sistem pendukung AI dari awal tanpa bergantung pada layanan cloud, API eksternal, atau model pre-trained yang sudah ada (seperti OpenAI, Claude, Llama, dll). 
+Tujuan akhirnya adalah memiliki AI Assistant pribadi yang berjalan sepenuhnya lokal di komputer pengguna.
+
+## Requirements
+- Windows 64-bit
+- Python 3.10+
+- Hardware minimal: CPU + RAM (Rekomendasi: NVIDIA GPU dengan CUDA support)
+
+## Setup Virtual Environment
+```powershell
+python -m venv venv
+.\venv\Scripts\activate
+```
+
+## Install Dependencies
+Project ini menggunakan `pyproject.toml` untuk manajemen dependensi minimal.
+```powershell
+pip install -e .[dev]
+```
+Dependensi utama:
+- PySide6 (Desktop GUI)
+- PyYAML (Configuration)
+- PyTorch (Hardware detection, Tensor/Autograd)
+
+## Menjalankan Aplikasi
+```powershell
+python run.py
+```
+
+## Menjalankan Test
+```powershell
+pytest tests/
+```

zyra_network-0.0.1/ai/__init__.py

@@ -0,0 +1 @@
+# AI Package

zyra_network-0.0.1/ai/agent/core_loop.py

@@ -0,0 +1,159 @@
+import logging
+import json
+import time
+from typing import List, Dict, Any, Callable
+from app.core.tools import execute_tool, TOOLS_SCHEMA
+
+class AgentState:
+    IDLE = "IDLE"
+    THINKING = "THINKING"
+    ACTING = "ACTING"
+    OBSERVING = "OBSERVING"
+    COMPLETED = "COMPLETED"
+    ERROR = "ERROR"
+
+class AutonomousAgent:
+    """
+    A robust ReAct-style Agent loop that can run long tasks and emit state changes.
+    This replaces the simple 5-iteration loop in the standard inference worker,
+    laying the groundwork for Proof-of-Useful-Work (PoUW).
+    """
+    def __init__(self, llm_client, max_iterations: int = 15):
+        self.llm_client = llm_client
+        self.max_iterations = max_iterations
+        self.state = AgentState.IDLE
+        self.logger = logging.getLogger("agent.core")
+        
+        # Callbacks for UI updates
+        self.on_state_change: Callable[[str], None] = None
+        self.on_log: Callable[[str], None] = None
+        
+        self.is_interrupted = False
+
+    def interrupt(self):
+        self.is_interrupted = True
+        self.llm_client.interrupt()
+
+    def _set_state(self, new_state: str):
+        self.state = new_state
+        if self.on_state_change:
+            self.on_state_change(self.state)
+
+    def _log(self, msg: str):
+        self.logger.info(msg)
+        if self.on_log:
+            self.on_log(msg)
+
+    def run_task(
+        self, 
+        task_prompt: str, 
+        system_context: str = "",
+        security_callback: Callable = None
+    ) -> str:
+        """
+        Executes a task autonomously until completion or max iterations.
+        """
+        self.is_interrupted = False
+        self._set_state(AgentState.THINKING)
+        self._log(f"Task started: {task_prompt}")
+        
+        messages = [
+            {"role": "system", "content": system_context},
+            {"role": "user", "content": task_prompt}
+        ]
+        
+        iteration = 0
+        final_answer = ""
+        
+        while iteration < self.max_iterations:
+            if self.is_interrupted:
+                self._log("Task interrupted by user.")
+                self._set_state(AgentState.ERROR)
+                return "Interrupted."
+                
+            iteration += 1
+            self._log(f"--- Iteration {iteration}/{self.max_iterations} ---")
+            self._set_state(AgentState.THINKING)
+            
+            try:
+                # Call LLM via standard API
+                response = self.llm_client.client.chat.completions.create(
+                    model=self.llm_client.model_name,
+                    messages=messages,
+                    tools=TOOLS_SCHEMA,
+                    temperature=0.7,
+                )
+                
+                message = response.choices[0].message
+                
+                # Log thought process if any
+                if message.content:
+                    self._log(f"Agent Thought:\n{message.content}")
+                    final_answer = message.content # Keep updating final answer
+                    
+                # Check for tool calls
+                if message.tool_calls:
+                    self._set_state(AgentState.ACTING)
+                    
+                    # Append assistant's tool calls to context
+                    assistant_msg = {
+                        "role": "assistant",
+                        "content": message.content,
+                        "tool_calls": [
+                            {
+                                "id": tc.id,
+                                "type": "function",
+                                "function": {
+                                    "name": tc.function.name,
+                                    "arguments": tc.function.arguments
+                                }
+                            } for tc in message.tool_calls
+                        ]
+                    }
+                    messages.append(assistant_msg)
+                    
+                    for tc in message.tool_calls:
+                        func_name = tc.function.name
+                        try:
+                            args = json.loads(tc.function.arguments)
+                        except json.JSONDecodeError:
+                            args = {}
+                            
+                        self._log(f"Executing Tool: {func_name}")
+                        
+                        # Execute Tool
+                        result_str = execute_tool(func_name, args, security_callback)
+                        
+                        self._set_state(AgentState.OBSERVING)
+                        
+                        # Truncate output for log UI to avoid freezing
+                        log_out = result_str
+                        if len(log_out) > 500:
+                            log_out = log_out[:500] + "... (truncated)"
+                        self._log(f"Tool Output:\n{log_out}")
+                        
+                        # Append observation to context
+                        messages.append({
+                            "role": "tool",
+                            "tool_call_id": tc.id,
+                            "content": result_str
+                        })
+                        
+                    # Loop back to THINKING with the new context
+                    continue 
+                else:
+                    # No tool calls means the agent is done
+                    self._log("Task completed successfully.")
+                    self._set_state(AgentState.COMPLETED)
+                    # Append final answer to messages to save it in memory if needed
+                    messages.append({"role": "assistant", "content": final_answer})
+                    return final_answer
+                    
+            except Exception as e:
+                self._log(f"Error during iteration: {str(e)}")
+                self._set_state(AgentState.ERROR)
+                return f"Error: {str(e)}"
+                
+        self._log("Max iterations reached without completion.")
+        self._set_state(AgentState.ERROR)
+        return final_answer + "\n[Warning: Task terminated because max iterations were reached.]"

zyra_network-0.0.1/ai/blockchain/ledger.py

@@ -0,0 +1,168 @@
+import sqlite3
+import os
+import json
+import hashlib
+import time
+
+class ZyraLedger:
+    """
+    A local SQLite-based blockchain ledger for the ZYRA token ecosystem.
+    Records Proof of Useful Work (PoUW) minting events and P2P transactions.
+    """
+    def __init__(self, data_dir: str):
+        self.db_path = os.path.join(data_dir, "ledger.db")
+        os.makedirs(data_dir, exist_ok=True)
+        self._init_db()
+
+    def _init_db(self):
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        
+        # Blocks Table
+        c.execute('''
+            CREATE TABLE IF NOT EXISTS blocks (
+                height INTEGER PRIMARY KEY AUTOINCREMENT,
+                timestamp REAL,
+                prev_hash TEXT,
+                merkle_root TEXT,
+                hash TEXT UNIQUE
+            )
+        ''')
+        
+        # Transactions Table
+        c.execute('''
+            CREATE TABLE IF NOT EXISTS transactions (
+                txid TEXT PRIMARY KEY,
+                block_height INTEGER,
+                sender TEXT,
+                receiver TEXT,
+                amount REAL,
+                type TEXT, -- 'MINT' (PoUW) or 'TRANSFER'
+                signature TEXT,
+                timestamp REAL,
+                FOREIGN KEY(block_height) REFERENCES blocks(height)
+            )
+        ''')
+        
+        # Balances Table (Account based for simplicity in local env)
+        c.execute('''
+            CREATE TABLE IF NOT EXISTS balances (
+                address TEXT PRIMARY KEY,
+                balance REAL DEFAULT 0.0
+            )
+        ''')
+        
+        conn.commit()
+        
+        # Create Genesis Block if empty
+        c.execute("SELECT COUNT(*) FROM blocks")
+        if c.fetchone()[0] == 0:
+            self._create_genesis_block(c)
+            conn.commit()
+            
+        conn.close()
+
+    def _create_genesis_block(self, cursor):
+        genesis_hash = hashlib.sha256(b"ZYRA_GENESIS_BLOCK").hexdigest()
+        cursor.execute('''
+            INSERT INTO blocks (timestamp, prev_hash, merkle_root, hash)
+            VALUES (?, ?, ?, ?)
+        ''', (time.time(), "0"*64, "0"*64, genesis_hash))
+
+    def get_balance(self, address: str) -> float:
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        c.execute("SELECT balance FROM balances WHERE address = ?", (address,))
+        row = c.fetchone()
+        conn.close()
+        return row[0] if row else 0.0
+
+    def add_pouw_reward(self, receiver_address: str, amount: float, task_proof: dict) -> str:
+        """
+        Mints new ZYRA tokens to the receiver for completing a valid AI task.
+        """
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        
+        timestamp = time.time()
+        
+        # Simple Block Generation (1 tx per block for now)
+        c.execute("SELECT hash, height FROM blocks ORDER BY height DESC LIMIT 1")
+        prev_block = c.fetchone()
+        prev_hash = prev_block[0]
+        new_height = prev_block[1] + 1
+        
+        tx_data = f"MINT:{receiver_address}:{amount}:{timestamp}:{json.dumps(task_proof)}"
+        txid = hashlib.sha256(tx_data.encode()).hexdigest()
+        
+        # Create Block
+        block_hash = hashlib.sha256(f"{new_height}{timestamp}{prev_hash}{txid}".encode()).hexdigest()
+        c.execute('''
+            INSERT INTO blocks (height, timestamp, prev_hash, merkle_root, hash)
+            VALUES (?, ?, ?, ?, ?)
+        ''', (new_height, timestamp, prev_hash, txid, block_hash))
+        
+        # Insert Transaction
+        c.execute('''
+            INSERT INTO transactions (txid, block_height, sender, receiver, amount, type, signature, timestamp)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+        ''', (txid, new_height, "SYSTEM", receiver_address, amount, "MINT", "POUW_VALIDATED", timestamp))
+        
+        # Update Balance
+        c.execute("INSERT OR IGNORE INTO balances (address, balance) VALUES (?, 0.0)", (receiver_address,))
+        c.execute("UPDATE balances SET balance = balance + ? WHERE address = ?", (amount, receiver_address))
+        
+        conn.commit()
+        conn.close()
+        return txid
+
+    def get_recent_transactions(self, limit=10):
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        c.execute('''
+            SELECT txid, type, amount, receiver, timestamp 
+            FROM transactions 
+            ORDER BY timestamp DESC LIMIT ?
+        ''', (limit,))
+        rows = c.fetchall()
+        conn.close()
+        
+        result = []
+        for r in rows:
+            result.append({
+                "txid": r[0],
+                "type": r[1],
+                "amount": r[2],
+                "receiver": r[3],
+                "timestamp": r[4]
+            })
+        return result
+
+    def add_withdraw_transaction(self, sender_address: str, amount: float, web3_txid: str):
+        conn = sqlite3.connect(self.db_path)
+        c = conn.cursor()
+        
+        timestamp = time.time()
+        c.execute("SELECT hash, height FROM blocks ORDER BY height DESC LIMIT 1")
+        prev_block = c.fetchone()
+        prev_hash = prev_block[0]
+        new_height = prev_block[1] + 1
+        
+        # We store the web3_txid as the local txid for cross-reference
+        txid = web3_txid if web3_txid else hashlib.sha256(f"WITHDRAW:{sender_address}:{amount}:{timestamp}".encode()).hexdigest()
+        
+        block_hash = hashlib.sha256(f"{new_height}{timestamp}{prev_hash}{txid}".encode()).hexdigest()
+        c.execute('''
+            INSERT INTO blocks (height, timestamp, prev_hash, merkle_root, hash)
+            VALUES (?, ?, ?, ?, ?)
+        ''', (new_height, timestamp, prev_hash, txid, block_hash))
+        
+        c.execute('''
+            INSERT INTO transactions (txid, block_height, sender, receiver, amount, type, signature, timestamp)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+        ''', (txid, new_height, sender_address, "WEB3_BRIDGE", amount, "WITHDRAW", "LOCAL_AUTH", timestamp))
+        
+        c.execute("UPDATE balances SET balance = balance - ? WHERE address = ?", (amount, sender_address))
+        
+        conn.commit()
+        conn.close()

zyra_network-0.0.1/ai/blockchain/pouw_validator.py

@@ -0,0 +1,72 @@
+import hashlib
+import json
+import time
+import math
+
+class PoUWValidator:
+    """
+    Proof of Useful Work (PoUW) Validator.
+    Instead of hashing empty blocks, we validate the computational work done during AI inference/agent tasks.
+    """
+    
+    BASE_REWARD_PER_TOKEN = 0.0001
+    DIFFICULTY_MULTIPLIER = 1.0
+
+    @classmethod
+    def calculate_reward(cls, tokens_generated: int, latency_ms: float, vram_mb: float, is_tool_call: bool) -> float:
+        """
+        Calculates the amount of ZYRA to reward based on the useful work performed.
+        Heavier tasks (more tokens, using tools, utilizing more VRAM) yield higher rewards.
+        """
+        if tokens_generated <= 0:
+            return 0.0
+            
+        # Base reward scales linearly with tokens generated
+        base = tokens_generated * cls.BASE_REWARD_PER_TOKEN
+        
+        # Multiplier for tool usage (Agentic work is more valuable than passive chat)
+        tool_multiplier = 1.5 if is_tool_call else 1.0
+        
+        # Hardware multiplier (Simulates higher rewards for offering more compute power)
+        # Using log10 so it doesn't scale infinitely
+        hw_multiplier = math.log10(max(10, vram_mb)) / math.log10(8192) if vram_mb > 0 else 1.0
+        
+        # Calculate final reward
+        reward = base * tool_multiplier * hw_multiplier * cls.DIFFICULTY_MULTIPLIER
+        
+        # Cap reward per single inference task to prevent abuse
+        return round(min(reward, 50.0), 6)
+
+    @classmethod
+    def generate_proof(cls, task_type: str, prompt: str, tokens: int, metrics: dict, wallet_address: str) -> dict:
+        """
+        Generates a cryptographic proof of the work done.
+        In a real P2P network, nodes would verify this proof against a model checkpoint.
+        For local simulation, we sign the metrics with the system timestamp.
+        """
+        timestamp = time.time()
+        
+        # Check if tools were used (indicated by task_type or metrics)
+        is_tool = task_type == 'AGENT_EXECUTION'
+        
+        vram_mb = metrics.get('vram_mb', 0.0)
+        latency_ms = metrics.get('latency_ms', 0.0)
+        
+        reward = cls.calculate_reward(tokens, latency_ms, vram_mb, is_tool)
+        
+        proof_payload = {
+            "task_type": task_type,
+            "wallet": wallet_address,
+            "tokens": tokens,
+            "metrics": metrics,
+            "timestamp": timestamp,
+            "reward": reward
+        }
+        
+        # Create a verifiable hash of the payload
+        payload_str = json.dumps(proof_payload, sort_keys=True)
+        proof_hash = hashlib.sha256(payload_str.encode()).hexdigest()
+        
+        proof_payload["proof_hash"] = proof_hash
+        
+        return proof_payload

zyra_network-0.0.1/ai/blockchain/wallet.py

@@ -0,0 +1,84 @@
+import os
+import json
+import hashlib
+import binascii
+try:
+    import ecdsa
+    HAS_ECDSA = True
+except ImportError:
+    HAS_ECDSA = False
+
+class ZyraWallet:
+    """
+    Handles cryptographic keys and signing for the local ZYRA token ecosystem.
+    Uses ECDSA SECP256k1 if available, otherwise falls back to a simulated hash (for testing).
+    """
+    def __init__(self, data_dir: str):
+        self.wallet_file = os.path.join(data_dir, "wallet.json")
+        self.private_key = None
+        self.public_key = None
+        self.address = None
+        self.load_or_create_wallet()
+
+    def _generate_simulated_keys(self):
+        """Fallback if ecdsa is not installed. Not secure!"""
+        import uuid
+        self.private_key = hashlib.sha256(str(uuid.uuid4()).encode()).hexdigest()
+        self.public_key = hashlib.sha256(self.private_key.encode()).hexdigest()
+        self.address = "Z" + hashlib.ripemd160(self.public_key.encode()).hexdigest() if hasattr(hashlib, 'ripemd160') else "Z" + self.public_key[:40]
+
+    def _generate_ecdsa_keys(self):
+        """Secure key generation using secp256k1 (Bitcoin standard)."""
+        sk = ecdsa.SigningKey.generate(curve=ecdsa.SECP256k1)
+        vk = sk.get_verifying_key()
+        
+        self.private_key = sk.to_string().hex()
+        self.public_key = vk.to_string().hex()
+        
+        # Hash pubkey to get address (SHA256 then RIPEMD160 usually, here we use double SHA256 for simplicity if RIPEMD missing)
+        pub_hash = hashlib.sha256(self.public_key.encode()).hexdigest()
+        self.address = "Z" + pub_hash[:40] # Z-prefix for ZYRA
+
+    def load_or_create_wallet(self):
+        if os.path.exists(self.wallet_file):
+            with open(self.wallet_file, 'r') as f:
+                data = json.load(f)
+                self.private_key = data.get('private_key')
+                self.public_key = data.get('public_key')
+                self.address = data.get('address')
+        else:
+            if HAS_ECDSA:
+                self._generate_ecdsa_keys()
+            else:
+                self._generate_simulated_keys()
+                
+            os.makedirs(os.path.dirname(self.wallet_file), exist_ok=True)
+            with open(self.wallet_file, 'w') as f:
+                json.dump({
+                    "private_key": self.private_key,
+                    "public_key": self.public_key,
+                    "address": self.address
+                }, f, indent=4)
+
+    def sign_transaction(self, tx_data: str) -> str:
+        """Signs a transaction payload string with the private key."""
+        if HAS_ECDSA and len(self.private_key) == 64: # 32 bytes hex
+            try:
+                sk = ecdsa.SigningKey.from_string(bytes.fromhex(self.private_key), curve=ecdsa.SECP256k1)
+                signature = sk.sign(tx_data.encode())
+                return signature.hex()
+            except Exception:
+                pass
+                
+        # Fallback simulated signature
+        return hashlib.sha256((self.private_key + tx_data).encode()).hexdigest()
+
+    @staticmethod
+    def verify_signature(public_key: str, signature: str, tx_data: str) -> bool:
+        if HAS_ECDSA and len(public_key) == 128:
+            try:
+                vk = ecdsa.VerifyingKey.from_string(bytes.fromhex(public_key), curve=ecdsa.SECP256k1)
+                return vk.verify(bytes.fromhex(signature), tx_data.encode())
+            except Exception:
+                return False
+        return True # Fallback mode accepts all

zyra_network-0.0.1/ai/brain/__init__.py

@@ -0,0 +1,6 @@
+from .base_model import BaseLanguageModel
+from .baseline_mlp import FixedContextMLP
+from .transformer import MyAIDecoderTransformer
+from .model_factory import create_model
+
+__all__ = ["BaseLanguageModel", "FixedContextMLP", "MyAIDecoderTransformer", "create_model"]

zyra_network-0.0.1/ai/brain/attention.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Any, Optional
+
+from .rope import RotaryPositionalEmbedding
+
+class MultiHeadAttention(nn.Module):
+    """
+    Multi-Head Causal Self-Attention.
+    Uses scaled dot-product attention with causal masking.
+    """
+    def __init__(self, config: Dict[str, Any]):
+        super().__init__()
+        self.hidden_size = config.get("hidden_size", 512)
+        self.num_heads = config.get("num_attention_heads", 8)
+        self.head_dim = self.hidden_size // self.num_heads
+        self.dropout = config.get("attention_dropout", 0.0)
+        self.bias = config.get("bias", False)
+        
+        if self.head_dim * self.num_heads != self.hidden_size:
+            raise ValueError(f"hidden_size ({self.hidden_size}) must be divisible by num_heads ({self.num_heads})")
+            
+        # Linear projections
+        # We use separate linears for clarity, though combined QKV is slightly faster.
+        # Combined is fine, let's use separate for clean RoPE application.
+        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.bias)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=self.bias)
+        
+        self.attn_dropout = nn.Dropout(self.dropout)
+        
+    def forward(self, x: torch.Tensor, rope: RotaryPositionalEmbedding) -> torch.Tensor:
+        """
+        Forward pass.
+        Args:
+            x: (B, T, D)
+            rope: RotaryPositionalEmbedding instance.
+        """
+        B, T, D = x.size()
+        
+        # 1. Projections
+        q = self.q_proj(x) # (B, T, D)
+        k = self.k_proj(x) # (B, T, D)
+        v = self.v_proj(x) # (B, T, D)
+        
+        # 2. Reshape to (B, H, T, HeadDim)
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        # 3. Apply RoPE to Query and Key
+        q, k = rope(q, k)
+        
+        # 4. Scaled Dot-Product Attention
+        # scores: (B, H, T, T)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        
+        # 5. Causal Mask
+        # We need a mask of shape (T, T) where upper triangle is -inf
+        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=1)
+        # Replace True with -inf
+        # Numerical safe masking
+        scores = scores.masked_fill(mask, float('-inf'))
+        
+        # 6. Softmax
+        # fp32 is safer for softmax
+        probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)
+        probs = self.attn_dropout(probs)
+        
+        # 7. Weighted Sum
+        # out: (B, H, T, HeadDim)
+        out = torch.matmul(probs, v)
+        
+        # 8. Merge Heads
+        # transpose back to (B, T, H, HeadDim) -> (B, T, D)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        
+        # 9. Output Projection
+        return self.o_proj(out)