1  use std::tensor::{Tensor, SparseMask, RotaryEmbed};
 2  use std::ops::{einsum, softmax, dropout, rms_norm};
 3  use std::distribute::Mesh;
 4  
 5  @kernel(grid=[N/256], block=[256])
 6  @autodiff(mode=reverse, order=2)
 7  @quantize(precision=fp8_e4m3, fallback=bf16)
 8  module TransformerBlock<T: TensorKind, D: Dim> {
 9  
10    // Flash attention w/ rotary position embeddings
11    @fuse(strategy=aggressive, mem_bound=shared)
12    fn multi_head_attn(
13      Q: Tensor<T, [B, H, S, D]>,
14      K: Tensor<T, [B, H, S, D]>,
15      V: Tensor<T, [B, H, S, D]>,
16      mask: SparseMask<[S, S], layout=causal>,
17      rope: RotaryEmbed<D, base=500000>,
18    ) -> Tensor<T, [B, H, S, D]> @where D % 128 == 0 {
19  
20      let Q_rot = rope.apply(Q, dim=-1);
21      let K_rot = rope.apply(K, dim=-1);
22  
23      let scale = 1.0 / D.sqrt().cast<T>();
24      let attn = einsum("bhsd,bhtd->bhst", Q_rot, K_rot)
25        |> .mul(scale)
26        |> mask.apply(_, fill=-inf)
27        |> softmax(dim=-1, algo=online_safe)
28        |> dropout(p=0.0, @compile_only);
29  
30      let out = einsum("bhst,bhtd->bhsd", attn, V);
31      return out |> rearrange("b h s d -> b s (h d)")
32                 |> linear(self.Wo, bias=none);
33    }
34  
35    // Gated FFN with SwiGLU activation
36    @pipeline(stages=4, interleave=1F1B)
37    @checkpoint(policy=selective, granularity=op)
38    fn feedforward(
39      x: Tensor<T, [B, S, E]>,
40    ) -> Tensor<T, [B, S, E]> {
41      let gate = x |> linear(self.Wg) |> silu();
42      let up   = x |> linear(self.Wu);
43      return (gate * up)
44        |> linear(self.Wd)
45        |> rms_norm(self.γ, eps=1e-6);
46    }
47  
48    // Forward pass w/ residual + MoE routing
49    @distribute(mesh=[dp=8, tp=4, pp=2, ep=16])
50    fn forward(x: Tensor<T, [B, S, E]>) -> Loss {
51      let residual = x;
52      let h = x |> rms_norm(self.γ1, eps=1e-6)
53                |> self.multi_head_attn(_, _, _)
54                |> residual.add(_);
55      let out = h |> rms_norm(self.γ2, eps=1e-6)
56                  |> self.feedforward(_)
57                  |> h.add(_);
58      return out |> sparse_moe(
59        experts=self.experts[0..E],
60        router=self.gate,
61        top_k=2,
62        capacity_factor=1.25,
63        @load_balance(aux_loss=0.01)
64      );
65    }
66  }