diff --git a/examples/mha/mha_module.py b/examples/mha/mha_module.py new file mode 100644 index 00000000..90b96bf1 --- /dev/null +++ b/examples/mha/mha_module.py @@ -0,0 +1,152 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Multi-Head Attention as an ``ark.Module`` subclass. + +Demonstrates: + - Subclassing ``ark.Module`` with ``ark.parameter()``. + - Using the functional API (``ark.matmul``, ``ark.reshape``, + ``ark.transpose``, ``ark.softmax``, ``ark.mul``). + - Numerical validation against a manual PyTorch reference implementation. + +Run: ``python examples/mha/mha_module.py`` +""" + +import math +import numpy as np +import torch +import ark + +# ---------- hyperparameters ---------- +batch_size = 1 +seq_len = 64 +d_model = 128 +n_heads = 4 + + +# ---------- ARK Module ---------- +class MultiHeadAttention(ark.Module): + """Scaled dot-product multi-head attention (no bias, no mask).""" + + def __init__( + self, d_model: int, n_heads: int, batch_size: int, seq_len: int + ): + super().__init__() + self.n_heads = n_heads + self.d_k = d_model // n_heads + self.scale = 1.0 / math.sqrt(self.d_k) + self.batch_size = batch_size + self.seq_len = seq_len + + # Projection weights: [d_model, d_model] + self.wq = ark.parameter([d_model, d_model], ark.fp16) + self.wk = ark.parameter([d_model, d_model], ark.fp16) + self.wv = ark.parameter([d_model, d_model], ark.fp16) + self.wo = ark.parameter([d_model, d_model], ark.fp16) + + def forward(self, x): + # x: [batch_size, seq_len, d_model] + # Shape is fixed at graph-build time; stored as instance attributes. + H = self.n_heads + dk = self.d_k + + # Linear projections + q = ark.matmul(x, self.wq) # [batch_size, seq_len, d_model] + k = ark.matmul(x, self.wk) + v = ark.matmul(x, self.wv) + + # Reshape to [batch_size, seq_len, H, dk] then transpose to [batch_size, H, seq_len, dk] + q = ark.transpose( + ark.reshape(q, [self.batch_size, self.seq_len, H, dk]), [0, 2, 1, 3] + ) + k = ark.transpose( + ark.reshape(k, [self.batch_size, self.seq_len, H, dk]), [0, 2, 1, 3] + ) + v = ark.transpose( + ark.reshape(v, [self.batch_size, self.seq_len, H, dk]), [0, 2, 1, 3] + ) + + # Scaled dot-product attention + # scores: [batch_size, H, seq_len, seq_len] + scores = ark.matmul(q, k, transpose_other=True) + scores = ark.mul(scores, self.scale) + attn = ark.softmax(scores) # along last axis + # context: [batch_size, H, seq_len, dk] + context = ark.matmul(attn, v) + + # Transpose back and reshape: [batch_size, seq_len, d_model] + context = ark.reshape( + ark.transpose(context, [0, 2, 1, 3]), + [self.batch_size, self.seq_len, H * dk], + ) + + # Output projection + out = ark.matmul(context, self.wo) + return out + + +# ---------- PyTorch reference ---------- +def pytorch_mha(x_np, wq, wk, wv, wo): + """Manual MHA in PyTorch matching the ARK implementation above.""" + x = torch.from_numpy(x_np).cuda() + B, S, D = x.shape + H = n_heads + dk = d_model // n_heads + + q = x @ torch.from_numpy(wq).cuda() + k = x @ torch.from_numpy(wk).cuda() + v = x @ torch.from_numpy(wv).cuda() + + q = q.view(B, S, H, dk).permute(0, 2, 1, 3) + k = k.view(B, S, H, dk).permute(0, 2, 1, 3) + v = v.view(B, S, H, dk).permute(0, 2, 1, 3) + + scores = q @ k.transpose(-2, -1) / math.sqrt(dk) + attn = torch.softmax(scores, dim=-1) + context = attn @ v + + context = context.permute(0, 2, 1, 3).contiguous().view(B, S, D) + out = context @ torch.from_numpy(wo).cuda() + return out.cpu().numpy() + + +# ---------- main ---------- +def main(): + ark.init() + + # Build graph + x_ark = ark.tensor([batch_size, seq_len, d_model], ark.fp16) + model = MultiHeadAttention(d_model, n_heads, batch_size, seq_len) + y_ark = model(x_ark) + + # Launch runtime + runtime = ark.Runtime() + runtime.launch() + + # Random inputs and weights (small magnitude for fp16 stability) + rng = np.random.RandomState(42) + x_np = (rng.randn(batch_size, seq_len, d_model) * 0.02).astype(np.float16) + wq_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16) + wk_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16) + wv_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16) + wo_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16) + + x_ark.from_numpy(x_np) + model.load_state_dict({"wq": wq_np, "wk": wk_np, "wv": wv_np, "wo": wo_np}) + + # Run ARK + runtime.run() + y_host = y_ark.to_numpy() + + # Run PyTorch reference + y_ref = pytorch_mha(x_np, wq_np, wk_np, wv_np, wo_np) + + # Validate + np.testing.assert_allclose(y_host, y_ref, atol=1e-2, rtol=1e-2) + max_err = np.max(np.abs(y_host - y_ref)) + print(f"MHA module test passed (max error: {max_err:.6f})") + + +if __name__ == "__main__": + main() diff --git a/examples/tutorial/tensor_ops_tutorial.py b/examples/tutorial/tensor_ops_tutorial.py new file mode 100644 index 00000000..805d2407 --- /dev/null +++ b/examples/tutorial/tensor_ops_tutorial.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Tensor ops tutorial — covers tensor creation, elementwise ops, +reductions, and ``Runtime`` usage with the current ARK functional API. + +Complements ``quickstart_tutorial.py`` by exercising more ops: + - ``ark.tensor``, ``ark.ones``, ``ark.zeros``, ``ark.constant`` + - ``ark.add``, ``ark.sub``, ``ark.mul``, ``ark.div`` + - ``ark.exp``, ``ark.sqrt``, ``ark.rsqrt`` + - ``ark.relu``, ``ark.gelu``, ``ark.sigmoid`` + - ``ark.reduce_sum``, ``ark.reduce_mean``, ``ark.reduce_max`` + - ``ark.layernorm`` + - ``ark.reshape``, ``ark.transpose`` + +Each op result is validated against a NumPy reference. + +Run: ``python examples/tutorial/tensor_ops_tutorial.py`` +""" + +import math + +import numpy as np +import ark + + +def tensor_ops_tutorial(): + ark.init() + + M, N = 64, 128 + + # ---- Tensor creation ---- + x = ark.tensor([M, N], ark.fp32) + y = ark.tensor([M, N], ark.fp32) + one = ark.ones([M, N], ark.fp32) + zero = ark.zeros([M, N], ark.fp32) + c5 = ark.constant(5.0, [M, N], ark.fp32) + + # ---- Elementwise arithmetic ---- + a_add = ark.add(x, y) + a_sub = ark.sub(x, y) + a_mul = ark.mul(x, y) + a_div = ark.div(x, ark.add(y, one)) # avoid div-by-zero + + # ---- Unary math ---- + a_exp = ark.exp(x) + a_sqrt = ark.sqrt(ark.add(x, c5)) # shift to positive + a_rsqrt = ark.rsqrt(ark.add(x, c5)) + + # ---- Activations ---- + a_relu = ark.relu(x) + a_gelu = ark.gelu(x) + a_sig = ark.sigmoid(x) + + # ---- Reductions ---- + r_sum = ark.reduce_sum(x, axis=-1) # [M, 1] + r_mean = ark.reduce_mean(x, axis=-1) + r_max = ark.reduce_max(x, axis=-1) + + # ---- Layernorm ---- + a_ln = ark.layernorm(x) + + # ---- Reshape / transpose ---- + a_reshape = ark.reshape(x, [M * N]) + a_trans = ark.transpose( + ark.reshape(x, [M, N // 2, 2]), [0, 2, 1] + ) # [M, 2, N//2] + + # ---- Launch runtime ---- + runtime = ark.Runtime() + runtime.launch() + + rng = np.random.RandomState(0) + x_np = (rng.randn(M, N) * 0.5).astype(np.float32) + y_np = (rng.randn(M, N) * 0.5).astype(np.float32) + + x.from_numpy(x_np) + y.from_numpy(y_np) + + runtime.run() + + # ---- Validate ---- + def check(name, ark_tensor, expected, atol=1e-5): + got = ark_tensor.to_numpy() + np.testing.assert_allclose( + got, expected, atol=atol, rtol=1e-4, err_msg=name + ) + + # Elementwise + check("add", a_add, x_np + y_np) + check("sub", a_sub, x_np - y_np) + check("mul", a_mul, x_np * y_np) + check("div", a_div, x_np / (y_np + 1.0)) + + # Unary + check("exp", a_exp, np.exp(x_np), atol=1e-4) + check("sqrt", a_sqrt, np.sqrt(x_np + 5.0), atol=1e-5) + check("rsqrt", a_rsqrt, 1.0 / np.sqrt(x_np + 5.0), atol=1e-5) + + # Activations + check("relu", a_relu, np.maximum(x_np, 0)) + + # GELU: approximate check (ARK uses erff-based GELU) + + gelu_ref = 0.5 * x_np * (1 + np.vectorize(math.erf)(x_np / np.sqrt(2))) + check("gelu", a_gelu, gelu_ref, atol=1e-4) + + sig_ref = 1.0 / (1.0 + np.exp(-x_np)) + check("sigmoid", a_sig, sig_ref, atol=1e-5) + + # Reductions (keepdims=True by default → [M, 1]) + check("reduce_sum", r_sum, x_np.sum(axis=-1, keepdims=True)) + check( + "reduce_mean", + r_mean, + x_np.mean(axis=-1, keepdims=True), + atol=1e-4, + ) + check("reduce_max", r_max, x_np.max(axis=-1, keepdims=True)) + + # Layernorm + mu = x_np.mean(axis=-1, keepdims=True) + # eps must match ark.layernorm default (1e-6) + var = ((x_np - mu) ** 2).mean(axis=-1, keepdims=True) + ln_ref = (x_np - mu) / np.sqrt(var + 1e-6) + check("layernorm", a_ln, ln_ref, atol=1e-4) + + # Reshape / transpose + check("reshape", a_reshape, x_np.reshape(M * N)) + check( + "transpose", + a_trans, + x_np.reshape(M, N // 2, 2).transpose(0, 2, 1), + ) + + # Constant tensors + check("ones", one, np.ones((M, N), dtype=np.float32)) + check("zeros", zero, np.zeros((M, N), dtype=np.float32)) + check("constant", c5, np.full((M, N), 5.0, dtype=np.float32)) + + print("Tensor ops tutorial passed — all ops validated!") + + +if __name__ == "__main__": + tensor_ops_tutorial() diff --git a/examples/tutorial/torch_interop_demo.py b/examples/tutorial/torch_interop_demo.py new file mode 100644 index 00000000..d8a1b50c --- /dev/null +++ b/examples/tutorial/torch_interop_demo.py @@ -0,0 +1,109 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +PyTorch interop demo — shows how ARK tensors and torch tensors +interact seamlessly. + +Demonstrates: + - ``Tensor.from_torch()`` — zero-copy ARK view of a CUDA torch tensor. + - ``Tensor.to_torch()`` — zero-copy torch view of an ARK tensor. + - Implicit torch→ARK conversion when passing torch tensors directly + to ``ark.*`` ops (no explicit conversion needed). The torch tensor + must be contiguous and on a CUDA device. + - ``Tensor.eval()`` — one-liner: build graph, run, return torch + tensor. + +Run: ``python examples/tutorial/torch_interop_demo.py`` +""" + +import numpy as np +import torch +import ark + + +def demo_from_to_torch(): + """Round-trip: torch → ARK → torch preserves data.""" + ark.init() + + t = torch.randn(4, 64, dtype=torch.float16, device="cuda") + + # torch → ARK (zero-copy; shares memory) + a = ark.Tensor.from_torch(t) + assert a.shape() == list(t.shape), "Shape mismatch" + + # Build a trivial identity graph so the runtime has work to do + out = ark.add(a, 0.0) + + runtime = ark.Runtime() + runtime.launch() + runtime.run() + + # ARK → torch (zero-copy via DLPack) + t2 = out.to_torch() + assert t2.shape == t.shape, "Shape mismatch after round-trip" + + # Adding 0.0 to fp16 values is exact; atol=0 is intentional. + np.testing.assert_allclose( + t2.cpu().numpy(), + t.cpu().numpy(), + atol=0, + err_msg="Round-trip data mismatch", + ) + print("[from/to_torch] Round-trip passed.") + + +def demo_implicit_conversion(): + """Pass torch tensors directly to ark ops (implicit conversion).""" + ark.init() + + x = torch.randn(8, 32, dtype=torch.float16, device="cuda") + y = torch.randn(8, 32, dtype=torch.float16, device="cuda") + + # ark.add accepts torch.Tensor inputs — implicit conversion + z = ark.add(x, y) + + runtime = ark.Runtime() + runtime.launch() + runtime.run() + + result = z.to_torch() + expected = x + y + np.testing.assert_allclose( + result.cpu().numpy(), + expected.cpu().numpy(), + atol=1e-3, + err_msg="Implicit conversion result mismatch", + ) + print("[implicit conv] torch tensors accepted by ark.add — passed.") + + +def demo_eval(): + """Tensor.eval() compiles and runs the graph in one call.""" + ark.init() + + x = torch.randn(4, 64, dtype=torch.float16, device="cuda") + + # eval() returns a torch.Tensor directly — no manual Runtime needed + result = ark.relu(x).eval() + + assert isinstance(result, torch.Tensor), "eval() should return torch.Tensor" + expected = torch.relu(x) + np.testing.assert_allclose( + result.cpu().numpy(), + expected.cpu().numpy(), + atol=1e-3, + err_msg="eval() result mismatch", + ) + print("[eval] One-liner eval passed.") + + +def main(): + demo_from_to_torch() + demo_implicit_conversion() + demo_eval() + print("\nAll torch interop demos passed!") + + +if __name__ == "__main__": + main()