Add gated MLP inference coverage

README.md

@@ -188,6 +188,11 @@ The initial implementation supports:
   `nn.Linear`, `nn.ReLU`, and `nn.Linear`, with matrix or contiguous
   `[batch, sequence, hidden]` inputs, on the ReLU-enabled runtime backends
   listed above
+- end-to-end `float32` inference for gated MLP blocks composed of parallel,
+  bias-free `nn.Linear` gate and up projections, `torch.nn.functional.silu`,
+  tensor multiplication, and a bias-free `nn.Linear` down projection, with
+  contiguous `[batch, sequence, hidden]` inputs, validated on the InfiniRT CPU
+  and NVIDIA backends
 
 The `torch.infini` module follows `torch.cuda` naming and semantics for the
 device and stream-management operations it implements. Stream priorities,

tests/test_gated_mlp.py

@@ -0,0 +1,70 @@
+import torch
+
+import torch_infini  # noqa: F401
+
+
+# Match PyTorch's TF32 test precision for CUDA matrix multiplication.
+TF32_TOLERANCE = {"rtol": 5e-3, "atol": 5e-3}
+
+
+class _GatedMLP(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.gate_proj = torch.nn.Linear(4, 6, bias=False)
+        self.up_proj = torch.nn.Linear(4, 6, bias=False)
+        self.down_proj = torch.nn.Linear(6, 4, bias=False)
+
+    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        gate = torch.nn.functional.silu(self.gate_proj(input_tensor))
+        return self.down_proj(gate * self.up_proj(input_tensor))
+
+
+def _copy_to_cpu(tensor: torch.Tensor) -> torch.Tensor:
+    result = torch.empty(tensor.shape, dtype=tensor.dtype)
+    result.copy_(tensor)
+    return result
+
+
+def _make_gated_mlp() -> _GatedMLP:
+    model = _GatedMLP().eval()
+    with torch.no_grad():
+        model.gate_proj.weight.copy_(
+            torch.linspace(-1.0, 1.0, steps=24, dtype=torch.float32).reshape(6, 4)
+        )
+        model.up_proj.weight.copy_(
+            torch.linspace(1.5, -0.5, steps=24, dtype=torch.float32).reshape(6, 4)
+        )
+        model.down_proj.weight.copy_(
+            torch.linspace(-0.75, 1.25, steps=24, dtype=torch.float32).reshape(4, 6)
+        )
+    return model
+
+
+def test_gated_mlp_inference_matches_cpu() -> None:
+    model = _make_gated_mlp()
+    input_cpu = torch.linspace(-2.0, 3.0, steps=24, dtype=torch.float32).reshape(
+        2, 3, 4
+    )
+    with torch.no_grad():
+        expected = model(input_cpu)
+
+    model = model.to("infini")
+    input_infini = input_cpu.to("infini")
+    with torch.no_grad():
+        result = model(input_infini)
+
+    assert not model.training
+    assert all(
+        module.bias is None
+        for module in model.modules()
+        if isinstance(module, torch.nn.Linear)
+    )
+    assert all(parameter.device.type == "infini" for parameter in model.parameters())
+    assert result.shape == input_cpu.shape
+    assert result.device.type == "infini"
+    assert not result.requires_grad
+    torch.testing.assert_close(
+        _copy_to_cpu(result),
+        expected,
+        **TF32_TOLERANCE,
+    )