Expand gemma_rmsnorm input dimension

src/sycl/Norm.h

@@ -21,7 +21,7 @@ inline std::tuple<int64_t, int64_t> _check_layer_norm_inputs(
     std::optional<torch::Tensor>& weight /* optional */,
     std::optional<torch::Tensor>& bias /* optional */) {
   CHECK_LAST_DIM_CONTIGUOUS(input);
-  TORCH_CHECK(input.dim() == 2 || input.dim() == 3, "input must be a 2D or 3D tensor");
+  TORCH_CHECK(input.dim() == 2 || input.dim() == 3 || input.dim() == 4, "input must be a 2D, 3D, or 4D tensor");
 #define TENSOR_CHECK(T)                          \
   if (T.has_value()) {                           \
     CHECK_LAST_DIM_CONTIGUOUS(T.value());        \

src/sycl/RMSNorm.cpp

@@ -48,15 +48,25 @@ struct RowStrides {
 };
 
 static inline RowStrides get_row_strides(const Tensor& t) {
-  TORCH_CHECK(t.dim() == 2 || t.dim() == 3, "get_row_strides: expected a 2D or 3D tensor, got ", t.dim(), "D");
+  TORCH_CHECK(
+      t.dim() == 2 || t.dim() == 3 || t.dim() == 4, "get_row_strides: expected a 2D/3D/4D tensor, got ", t.dim(), "D");
   if (t.dim() == 2) {
     return {t.stride(0), 1, 0};
   }
-  // 3D
-  int64_t outer_stride = t.stride(0);
-  int64_t inner_size = t.size(1);
-  int64_t inner_stride = t.stride(1);
-  if (t.size(0) == 1 || outer_stride == inner_size * inner_stride) {
+  if (t.dim() == 4) {
+    // 4D only: the leading batch-like dimension (dim 0) must be size 1,
+    // since our two-level (outer, inner) stride formula cannot represent
+    // a third level of striding.
+    TORCH_CHECK(
+        t.size(0) == 1, "get_row_strides: leading dimension 0 must have size 1 for a 4D tensor, got size ", t.size(0));
+  }
+
+  // For 3D/4D tensors, the outer is the second-to-last dimension and
+  // the inner is the last dimension.
+  int64_t outer_stride = t.stride(-3);
+  int64_t inner_size = t.size(-2);
+  int64_t inner_stride = t.stride(-2);
+  if (t.size(-3) == 1 || outer_stride == inner_size * inner_stride) {
     // Flattenable: a single stride describes all rows.
     return {inner_stride, 1, 0};
   }

tests/test_norm.py

@@ -440,6 +440,75 @@ def test_gemma_norm_3d_non_flattenable(
     torch.testing.assert_close(y_ref, y, rtol=1e-3, atol=1e-3)
 
 
+###############################################################################
+# 4D tensor tests for gemma_rmsnorm
+###############################################################################
+
+
+def _make_non_flattenable_4d(num_tokens, num_heads, head_dim, dtype, extra_heads=4):
+    """Create a 4D tensor [1, tokens, heads, head_dim] whose row strides are
+    not flattenable by a single outer stride.
+    """
+    total_heads = num_heads + extra_heads
+    full = torch.randn(
+        1, num_tokens, total_heads * head_dim, device=device, dtype=dtype
+    )
+    q_flat = full[:, :, : num_heads * head_dim]
+    q_4d = q_flat.unflatten(-1, (num_heads, head_dim))
+    assert q_4d.size(0) == 1
+    assert q_4d.stride(-3) == total_heads * head_dim
+    assert q_4d.stride(-3) != q_4d.size(-2) * q_4d.stride(-2)
+    return q_4d
+
+
+@pytest.mark.parametrize("num_tokens", [1, 7])
+@pytest.mark.parametrize("num_heads", [4, 8])
+@pytest.mark.parametrize("head_dim", [64, 128])
+@pytest.mark.parametrize("dtype", [torch.float16])
+@pytest.mark.parametrize("specify_out", [True, False])
+def test_gemma_norm_4d(num_tokens, num_heads, head_dim, dtype, specify_out):
+    x = torch.randn(1, num_tokens, num_heads, head_dim, device=device, dtype=dtype)
+    w = torch.randn(head_dim, device=device, dtype=dtype)
+
+    y_ref = gemma_rms_norm(x, w)
+    if specify_out:
+        y = torch.empty_like(x)
+        sgl_kernel.gemma_rmsnorm(x, w, out=y)
+    else:
+        y = sgl_kernel.gemma_rmsnorm(x, w)
+
+    torch.testing.assert_close(y_ref, y, **norm_tolerances(dtype))
+
+
+@pytest.mark.parametrize("num_tokens", [7, 32])
+@pytest.mark.parametrize("num_heads", [4, 8])
+@pytest.mark.parametrize("head_dim", [64, 128])
+@pytest.mark.parametrize("dtype", [torch.float16])
+@pytest.mark.parametrize("specify_out", [True, False])
+def test_gemma_norm_4d_non_flattenable(
+    num_tokens, num_heads, head_dim, dtype, specify_out
+):
+    x = _make_non_flattenable_4d(num_tokens, num_heads, head_dim, dtype)
+    w = torch.randn(head_dim, device=device, dtype=dtype)
+
+    y_ref = gemma_rms_norm(x.clone(), w)
+    if specify_out:
+        y = torch.empty_strided(x.shape, x.stride(), device=device, dtype=x.dtype)
+        sgl_kernel.gemma_rmsnorm(x, w, out=y)
+    else:
+        y = sgl_kernel.gemma_rmsnorm(x, w)
+
+    torch.testing.assert_close(y_ref, y, **norm_tolerances(dtype))
+
+
+def test_gemma_norm_4d_invalid_leading_dim_size_raises():
+    x = torch.randn(2, 7, 4, 128, device=device, dtype=torch.float16)
+    w = torch.randn(128, device=device, dtype=torch.float16)
+
+    with pytest.raises(RuntimeError, match="leading dimension 0 must have size 1"):
+        sgl_kernel.gemma_rmsnorm(x, w)
+
+
 ###############################################################################
 # Mixed input/weight dtype tests
 ###############################################################################