LLM Preliminaries

# English comments are used as requested in user instructions.
import numpy as np


# Let's define the SiLU activation function, as it's central to the FFN block.
# SiLU(x) = x * sigmoid(x)
def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def silu(x):
    return x * sigmoid(x)


# --- 1. Setup Simulation Parameters ---

# Number of devices (e.g., GPUs) to simulate for tensor parallelism.
num_devices = 4

# Define the dimensions of the model and input.
# For demonstration, these values are smaller than in a real model like Llama.
batch_size = 8
hidden_dim = 1024  # Dimension of the input and output 'x'
intermediate_dim = 4096  # The expanded inner dimension of the FFN

# Ensure intermediate_dim is divisible by num_devices for clean splitting.
assert intermediate_dim % num_devices == 0, (
    "intermediate_dim must be divisible by num_devices"
)

print(f"Simulating FFN parallelization across {num_devices} devices.")
print(
    f"Dimensions: Batch={batch_size}, Hidden={hidden_dim}, Intermediate={intermediate_dim}\n"
)

# --- 2. Initialize Input and Weights (Ground Truth) ---

# Create a random input tensor 'x'.
# In a real scenario, this comes from the previous Attention Block.
np.random.seed(42)
x = np.random.randn(batch_size, hidden_dim)

# Initialize the three weight matrices for the entire FFN block.
# These are the "unsplit" weights that would exist on a single device.
W_gate = np.random.randn(hidden_dim, intermediate_dim)
W_up = np.random.randn(hidden_dim, intermediate_dim)
W_down = np.random.randn(intermediate_dim, hidden_dim)


# --- 3. Baseline: Single-Device FFN Calculation ---
# This is our ground truth to verify the parallelized version.

print("--- Running Baseline Calculation (Single Device) ---")

# The formula is: Output = W_down(SiLU(W_gate(x)) * W_up(x))
gate_proj = x @ W_gate
up_proj = x @ W_up
# Element-wise multiplication
fused_hidden = silu(gate_proj) * up_proj
single_device_output = fused_hidden @ W_down

print(f"Shape of single_device_output: {single_device_output.shape}\n")


# --- 4. Parallelized FFN Calculation ---
# This simulates the process across multiple devices.

print("--- Running Parallelized Calculation (Simulated Multi-Device) ---")

# Step 0: Split the weights across the devices.
# W_gate and W_up are split by COLUMNS (axis=1). This is "Column Parallelism".
W_gate_shards = np.array_split(W_gate, num_devices, axis=1)
W_up_shards = np.array_split(W_up, num_devices, axis=1)

# W_down is split by ROWS (axis=0). This is "Row Parallelism".
W_down_shards = np.array_split(W_down, num_devices, axis=0)

print(f"Original W_gate shape: {W_gate.shape}")
print(f"Split W_gate_shard[0] shape: {W_gate_shards[0].shape}")
print(f"Original W_down shape: {W_down.shape}")
print(f"Split W_down_shard[0] shape: {W_down_shards[0].shape}\n")


# This list will hold the partial output from each device before the final communication.
partial_outputs = []

# --- Loop to simulate computation on each device ---
for i in range(num_devices):
    print(f"-> Simulating computation on Device {i}...")

    # The full input 'x' is available on every device (broadcast).
    device_input = x

    # Get the weight shards for this specific device.
    w_gate_i = W_gate_shards[i]
    w_up_i = W_up_shards[i]
    w_down_i = W_down_shards[i]

    # Step 1 & 2: Column and Row Parallelism Calculation
    # Each device performs its computation independently. NO communication is needed yet.

    # Column Parallelism part:
    # Each device computes a slice of the gate and up projections.
    gate_proj_i = device_input @ w_gate_i
    up_proj_i = device_input @ w_up_i

    # Local activation and element-wise multiplication.
    fused_hidden_i = silu(gate_proj_i) * up_proj_i

    # Row Parallelism part:
    # The intermediate result is multiplied by the corresponding ROW-split part of W_down.
    # The result is a partial output that has the final correct dimension.
    output_i = fused_hidden_i @ w_down_i

    print(f"   Shape of partial output on Device {i}: {output_i.shape}")

    # Store the partial result.
    partial_outputs.append(output_i)

print("\n--- Step 3: All-Reduce Aggregation ---")

# In a real system, this would be a single, highly optimized MPI or NCCL call.
# The all_reduce operation sums the tensors from all devices element-wise
# and makes the final result available on all devices.
# Here, we simply sum the partial results from our list.
parallel_device_output = sum(partial_outputs)

print(f"Shape of aggregated parallel_device_output: {parallel_device_output.shape}\n")


# --- 5. Verification ---
# Compare the output from the single-device calculation with the parallelized one.

print("--- Verification ---")
are_close = np.allclose(single_device_output, parallel_device_output, atol=1e-6)
print(
    f"Do the results from single-device and parallelized computation match? -> {are_close}"
)

if are_close:
    print(
        "✅ Success! The parallel FFN implementation correctly reproduces the original output."
    )
else:
    print("❌ Failure! The outputs do not match.")

def linear(
    x: mx.array,
    w: mx.array,
    bias: mx.array | None = None,
) -> mx.array:
    assert x.shape[-1] == w.shape[-1], "x.shape[-1] != w.shape[-]"
    return x @ w.T + (bias if bias is not None else 0)

class MultiHeadAttention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()
        self.n_embd = args.n_embd
        self.n_head = args.n_head
        self.head_dim = self.n_embd // self.n_head

        self.scale = self.head_dim**-0.5

        # D x 3D (Q, K, V)
        self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=True)
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=True)

    def __call__(
        self,
        x: mx.array,
        mask: Optional[mx.array | Literal["causal"]] = None,
        cache: Optional[Any] = None,
    ) -> mx.array:
        N, L, E = x.shape
        H, D = self.n_head, self.head_dim

        assert E == H * D, (
            f"Embedding size {E} must be divisible by number of heads {H}."
        )

        qkv = self.c_attn(x)
        # N x L x E
        q, k, v = mx.split(qkv, 3, axis=-1)

        # N x L x (H * D)
        q = q.reshape(N, L, H, D)
        k = k.reshape(N, L, H, D)
        v = v.reshape(N, L, H, D)

        # multihead attention expects
        # N x H x L x D
        q = q.swapaxes(1, 2)
        k = k.swapaxes(1, 2)
        v = v.swapaxes(1, 2)

        if cache is not None:
            k, v = cache.update_and_fetch(k, v)

        # N x H x L x D
        x = scaled_dot_product_attention(
            q, k, v, scale=self.scale, cache=cache, mask=mask
        )
        # N x L x (H * D)
        x = x.swapaxes(1, 2)
        x = x.reshape(N, L, H * D)
        return self.c_proj(x)

uv venv
uv pip install openai

import json
import os

from openai import OpenAI

# --- 1. Setup (Modified Section) ---
# Read API key and optional Base URL from environment variables.
# This is a best practice for security and flexibility, allowing you
# to use different API providers or proxies without changing the code.
api_key = "xxxx"
base_url = "https://openrouter.ai/api/v1"

# Check if the API key is provided, which is essential.
if not api_key:
    print("Error: The OPENAI_API_KEY environment variable is not set.")
    exit()

# Initialize the OpenAI client.
# The `base_url` parameter is optional. If it's None (not set as an env var),
# the client will default to OpenAI's official API endpoint.
try:
    client = OpenAI(
        api_key=api_key,
        base_url=base_url,
    )
    print(f"--- Client initialized. Using Base URL: {client.base_url} ---")
except Exception as e:
    print(f"Error initializing OpenAI client: {e}")
    exit()


# --- 2. Define the Tool (Your Python Function) ---
# This part remains unchanged.
def get_current_weather(location, unit="celsius"):
    """
    Get the current weather in a given location.

    Args:
        location (str): The city and state, e.g., "San Francisco, CA".
        unit (str): The unit for the temperature, can be "celsius" or "fahrenheit".

    Returns:
        str: A JSON string with weather information.
    """
    print(f"--- Executing 'get_current_weather' for {location} ---")
    if "tokyo" in location.lower():
        return json.dumps(
            {
                "location": "Tokyo",
                "temperature": "15",
                "unit": unit,
                "forecast": "rainy",
            }
        )
    elif "san francisco" in location.lower():
        return json.dumps(
            {
                "location": "San Francisco",
                "temperature": "22",
                "unit": unit,
                "forecast": "sunny",
            }
        )
    elif "paris" in location.lower():
        return json.dumps(
            {
                "location": "Paris",
                "temperature": "18",
                "unit": unit,
                "forecast": "cloudy",
            }
        )
    else:
        return json.dumps({"location": location, "temperature": "unknown"})


# --- 3. Main Conversation Loop ---
# This part remains unchanged.
def run_conversation():
    messages = [
        {
            "role": "user",
            "content": "What's the weather like in San Francisco and Tokyo?",
        }
    ]

    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "Get the current weather in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g., San Francisco, CA",
                        },
                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
                    },
                    "required": ["location"],
                },
            },
        }
    ]

    print("\n--- Step 1: First API Call (Model decides to use a tool) ---")
    response = client.chat.completions.create(
        model="openai/gpt-4.1",
        messages=messages,
        tools=tools,
        tool_choice="auto",
    )

    response_message = response.choices[0].message
    print("\n[Model's First Response - Tool Call Request]")
    print(response_message)

    tool_calls = response_message.tool_calls
    if tool_calls:
        messages.append(response_message)
        available_functions = {"get_current_weather": get_current_weather}

        for tool_call in tool_calls:
            function_name = tool_call.function.name
            function_to_call = available_functions[function_name]
            function_args = json.loads(tool_call.function.arguments)
            function_response = function_to_call(**function_args)
            messages.append(
                {
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": function_response,
                }
            )

        print(
            "\n--- Step 2: Second API Call (Sending tool results back to the model) ---"
        )
        second_response = client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
        )

        final_response = second_response.choices[0].message
        print("\n--- Step 3: Final Answer from Model ---")
        print(final_response.content)
    else:
        print("\n--- Final Answer from Model (No Tool Call) ---")
        print(response_message.content)


# Run the main function
if __name__ == "__main__":
    run_conversation()

传统 Function Calling:
AI 应用 ←→ Function Handler

MCP 协议:
AI 应用 ←→ MCP 客户端 ←→ MCP 服务器