[KernelGen][Nvidia] Add dsplit operator with view implementation

benchmark/test_dsplit.py

@@ -0,0 +1,56 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base, consts
+
+
+def dsplit_input_fn(shape, dtype, device):
+    inp = base.generate_tensor_input(shape, dtype, device)
+    # Use integer split (equal chunks)
+    sections = 2
+    yield inp, sections
+    if base.Config.bench_level == consts.BenchLevel.COMPREHENSIVE:
+        yield inp, 4
+
+
+class DsplitBenchmark(base.GenericBenchmark):
+    def get_input_iter(self, cur_dtype) -> Generator:
+        shapes = [
+            (64, 128, 32),
+            (128, 256, 64),
+            (256, 512, 128),
+            (512, 1024, 256),
+        ]
+
+        for shape in shapes:
+            yield from self.input_fn(shape, cur_dtype, self.device)
+
+
+@pytest.mark.dsplit
+def test_perf_dsplit():
+    def dsplit_wrapper(input, sections):
+        return torch.ops.aten.dsplit.int(input, sections)
+
+    bench = DsplitBenchmark(
+        input_fn=dsplit_input_fn,
+        op_name="dsplit",
+        torch_op=dsplit_wrapper,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

conf/operators.yaml

@@ -2480,6 +2480,24 @@ ops:
       - Math
     stages:
       - alpha: '5.4'
+  - id: dsplit
+    description: |
+      Split a tensor along the third axis (depth-wise). Pure layout operation returning zero-copy views.
+    for:
+      - dsplit.int
+      - dsplit.array
+    marks:
+      - dsplit
+    labels:
+      - aten
+      - KernelGen
+      - skip_precision_check
+    kind:
+      - Tensor
+    stages:
+      - alpha: '5.4'
+    device:
+      - nvidia
   - id: dequantize
     description: >
       Returns an fp32 Tensor by dequantizing a quantized Tensor.

src/flag_gems/__init__.py

@@ -394,6 +394,8 @@ def torch_ge(v):
     ("deg2rad", deg2rad),
     ("deg2rad.out", deg2rad_out),
     ("deg2rad_", deg2rad_),
+    ("dsplit.int", dsplit),
+    ("dsplit.array", dsplit),
     ("dequantize", dequantize),
     ("dequantize.self", dequantize),
     ("diag", diag),

src/flag_gems/ops/__init__.py

@@ -259,6 +259,7 @@
 )
 from flag_gems.ops.dot import dot
 from flag_gems.ops.dropout import dropout, dropout_backward
+from flag_gems.ops.dsplit import dsplit
 from flag_gems.ops.elu import elu, elu_, elu_backward
 from flag_gems.ops.embedding import embedding, embedding_backward
 from flag_gems.ops.embedding_dense_backward import embedding_dense_backward
@@ -1018,6 +1019,7 @@
     "cumsum_out",
     "deg2rad",
     "deg2rad_",
+    "dsplit",
     "deg2rad_out",
     "dequantize",
     "diag",

src/flag_gems/ops/dsplit.py

@@ -0,0 +1,37 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import List, Union
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def dsplit(input: torch.Tensor, indices_or_sections: Union[int, List[int]]):
+    """Split a tensor along the third axis (depth-wise).
+
+    This is equivalent to torch.tensor_split with dim=2 for 3D+ tensors.
+    Returns a tuple of views (zero-copy).
+    """
+    logger.debug("GEMS DSPLIT")
+
+    # dsplit splits along dim=2 (depth)
+    if isinstance(indices_or_sections, int):
+        # Equal splits
+        return torch.split(input, input.shape[2] // indices_or_sections, dim=2)
+    else:
+        # Custom split indices
+        return torch.tensor_split(input, indices_or_sections, dim=2)

tests/test_dsplit.py

@@ -0,0 +1,65 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+DSPLIT_CONFIGS = [
+    # (shape, indices_or_sections)
+    # Integer splits (equal chunks)
+    ((4, 6, 8), 2),
+    ((8, 4, 12), 4),
+    ((12, 8, 16), 4),
+    ((16, 16, 8), 2),
+    ((20, 10, 15), 5),
+    # List splits (custom indices)
+    ((4, 6, 8), [2]),
+    ((4, 6, 8), [4]),
+    ((6, 8, 10), [2, 6]),
+    ((8, 4, 12), [3, 6, 9]),
+    ((10, 5, 15), [5, 10]),
+    # 4D tensors
+    ((4, 6, 8, 3), 2),
+    ((8, 4, 12, 5), 4),
+    ((6, 8, 10, 7), [2, 6]),
+    ((10, 5, 15, 3), [5, 10]),
+]
+
+
+@pytest.mark.dsplit
+@pytest.mark.parametrize("shape, indices_or_sections", DSPLIT_CONFIGS)
+def test_accuracy_dsplit(shape, indices_or_sections):
+    inp = torch.randn(shape, dtype=torch.float32, device=flag_gems.device)
+    ref_inp = utils.to_reference(inp, True)
+
+    if isinstance(indices_or_sections, int):
+        ref_out = torch.ops.aten.dsplit.int(ref_inp, indices_or_sections)
+    else:
+        ref_out = torch.ops.aten.dsplit.array(ref_inp, indices_or_sections)
+
+    with flag_gems.use_gems():
+        if isinstance(indices_or_sections, int):
+            res_out = torch.ops.aten.dsplit.int(inp, indices_or_sections)
+        else:
+            res_out = torch.ops.aten.dsplit.array(inp, indices_or_sections)
+
+    assert len(res_out) == len(
+        ref_out
+    ), f"Length mismatch: {len(res_out)} vs {len(ref_out)}"
+    for i, (res_chunk, ref_chunk) in enumerate(zip(res_out, ref_out)):
+        utils.gems_assert_close(res_chunk, ref_chunk, torch.float32)