Expand gemma_rmsnorm input dimension - #367
Conversation
0be1a4f to
0cd7d1f
Compare
There was a problem hiding this comment.
Pull request overview
This PR extends the SYCL RMSNorm stride-handling to accept Gemma-style 4D inputs (with a constrained leading dimension) and adds pytest coverage for 4D contiguous and non-flattenable view layouts, aligning sgl_kernel.gemma_rmsnorm with Gemma3 usage patterns in sglang.
Changes:
- Allow
_check_layer_norm_inputsto accept 4D inputs. - Extend
get_row_strides()to support 4D tensors (requiring dim0 size 1) and derive strides from the last 3 dimensions. - Add new 4D
gemma_rmsnormtest cases, including a non-flattenable QKV-slice-like layout and an invalid-leading-dim error test.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| tests/test_norm.py | Adds 4D test coverage for gemma_rmsnorm, including non-flattenable stride patterns and error handling. |
| src/sycl/RMSNorm.cpp | Extends row-stride derivation to 4D and uses it in gemma_rmsnorm/rmsnorm kernels. |
| src/sycl/Norm.h | Updates input-dimension validation to allow 4D inputs in the shared norm input checker. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // 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.dim() == 3) && (t.size(0) == 1)) || t.size(-2) == 1 || outer_stride == inner_size * inner_stride) { | ||
| // Flattenable: a single stride describes all rows. | ||
| return {inner_stride, 1, 0}; | ||
| } |
| // 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)); |
There was a problem hiding this comment.
Can we support cases of t.size(0) > 1 ?
There was a problem hiding this comment.
We don't support t.size(0) > 1 because the only 4d size of gemma3's q/k input is [1, h, s, head_dim]. We can support t.size(0) > 1 for extensibility.
| 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) { |
There was a problem hiding this comment.
Please be aware of these copilot comments and ensure the tests cover these corner cases.
According to gemma3 model in sglang,
gemma_rmsnormmay receive a 4d input. This PR expand the input dimension to flattenable/unflattenable 4d shape and add corresponding test.