diff --git a/src/sycl/Norm.h b/src/sycl/Norm.h index 6330c2d71..a9ff20787 100644 --- a/src/sycl/Norm.h +++ b/src/sycl/Norm.h @@ -21,7 +21,7 @@ inline std::tuple _check_layer_norm_inputs( std::optional& weight /* optional */, std::optional& 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()); \ diff --git a/src/sycl/RMSNorm.cpp b/src/sycl/RMSNorm.cpp index f414ad635..b0d2f2abe 100644 --- a/src/sycl/RMSNorm.cpp +++ b/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}; } diff --git a/tests/test_norm.py b/tests/test_norm.py index 7cf469a82..5e7810e9d 100644 --- a/tests/test_norm.py +++ b/tests/test_norm.py @@ -254,7 +254,10 @@ def test_fused_add_rmsnorm_3d(batch_size, seq_len, hidden_size, dtype): @pytest.mark.parametrize("batch_size", [1, 4, 19]) @pytest.mark.parametrize("seq_len", [1, 7, 32]) -@pytest.mark.parametrize("hidden_size", [111, 1024, 4096]) +# hidden_size=1 exercises the "other dim == 1" shape (excluding the leading +# batch dim) that bypasses the flattenable fast-path check but still yields +# a correct result. +@pytest.mark.parametrize("hidden_size", [1, 111, 1024, 4096]) @pytest.mark.parametrize("dtype", [torch.float16]) @pytest.mark.parametrize("specify_out", [True, False]) def test_gemma_norm_3d(batch_size, seq_len, hidden_size, dtype, specify_out): @@ -440,6 +443,140 @@ 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]) +# num_heads=1 and head_dim=1 exercise the "other dim == 1" shapes (excluding +# the leading batch/token dims) that bypass the flattenable fast-path check +# but still yield a correct result. +@pytest.mark.parametrize("num_heads", [1, 4, 8]) +@pytest.mark.parametrize("head_dim", [1, 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)) + + +############################################################################### +# Non-contiguous 4D tensor tests (sliced last-dim, flattenable leading dims) +############################################################################### + + +def _make_non_contiguous_4d(num_tokens, num_heads, head_dim, dtype, extra=64): + """Create a last-dim-non-contiguous 4D tensor [1, tokens, heads, head_dim] + by slicing a larger tensor along the last dimension.""" + full = torch.randn( + 1, num_tokens, num_heads, head_dim + extra, device=device, dtype=dtype + ) + view = full[..., :head_dim] + assert view.stride(-1) == 1 + assert view.stride(-2) == head_dim + extra + return view + + +@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]) +def test_gemma_norm_4d_non_contiguous(num_tokens, num_heads, head_dim, dtype): + x_nc = _make_non_contiguous_4d(num_tokens, num_heads, head_dim, dtype) + w = torch.randn(head_dim, device=device, dtype=dtype) + + y_ref = gemma_rms_norm(x_nc.clone(), w) + y = sgl_kernel.gemma_rmsnorm(x_nc, w) + + torch.testing.assert_close(y_ref, y, **norm_tolerances(dtype)) + + +def test_gemma_norm_4d_non_flattenable_row_strides(): + x = _make_non_flattenable_4d(7, 4, 128, torch.float16) + w = torch.randn(x.size(-1), device=device, dtype=x.dtype) + + assert x.size(-2) > 1 + assert x.stride(-2) != 0 + assert x.stride(-3) != x.size(-2) * x.stride(-2) + + y = torch.empty_strided(x.shape, x.stride(), device=device, dtype=x.dtype) + y_ref = gemma_rms_norm(x.clone(), w) + sgl_kernel.gemma_rmsnorm(x, w, out=y) + + torch.testing.assert_close(y_ref, y, **norm_tolerances(x.dtype)) + + +def test_gemma_norm_4d_non_flattenable_unaligned_row_strides(): + full = torch.randn(1, 7, 8, 130, device=device, dtype=torch.float16) + x = full[:, :, :4, :128] + w = torch.randn(x.size(-1), device=device, dtype=x.dtype) + + assert x.stride(-3) != x.size(-2) * x.stride(-2) + assert x.size(-1) % 8 == 0 + assert x.stride(-2) % 8 != 0 + + y = torch.empty_strided(x.shape, x.stride(), device=device, dtype=x.dtype) + y_ref = gemma_rms_norm(x.clone(), w) + sgl_kernel.gemma_rmsnorm(x, w, out=y) + + torch.testing.assert_close(y_ref, y, **norm_tolerances(x.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 ###############################################################################