fix: resolve merge conflict for _jagged_to_padded_dense_forward PR (flagos-ai#3809)

Rebase onto origin/master (6f0fa85): resolved ops/__init__.py import/__all__
conflict by rebuilding PR changes cleanly from origin/master. Added operator
entry in correct alphabetical position across all registration files.

Author: XDYuanzhuLee

benchmark/test_jagged_to_padded_dense_forward.py

@@ -0,0 +1,47 @@
+import numpy as np
+import pytest
+import torch
+
+from . import base, consts
+
+# Jagged to padded dense forward benchmark
+JAGGED_TO_PADDED_SHAPES = [
+    (8, 8),
+    (16, 16),
+    (32, 32),
+    (64, 64),
+    (128, 64),
+    (256, 128),
+    (512, 256),
+]
+
+
+class JaggedToPaddedDenseForwardBenchmark(base.Benchmark):
+    def set_shapes(self, shape_file_path=None):
+        self.shapes = JAGGED_TO_PADDED_SHAPES
+
+    def get_input_iter(self, cur_dtype):
+        for batch_size, max_length in self.shapes:
+            # Generate random sequence lengths
+            np.random.seed(42)
+            seq_lengths = np.random.randint(1, max_length + 1, size=batch_size).tolist()
+
+            # Create offsets tensor (cumulative)
+            offsets = [0] + list(np.cumsum(seq_lengths).astype(int).tolist())
+            offsets = torch.tensor(offsets, device=self.device, dtype=torch.int64)
+
+            # Create values tensor (concatenated sequences)
+            total_length = sum(seq_lengths)
+            values = torch.randn(total_length, dtype=cur_dtype, device=self.device)
+
+            yield values, [offsets], [max_length], 0.0
+
+
+@pytest.mark.jagged_to_padded_dense_forward
+def test_jagged_to_padded_dense_forward():
+    bench = JaggedToPaddedDenseForwardBenchmark(
+        op_name="jagged_to_padded_dense_forward",
+        torch_op=torch.ops.aten._jagged_to_padded_dense_forward,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

conf/operators.yaml

@@ -3754,6 +3754,18 @@ ops:
       - Math
     stages:
       - stable: '5.3'
+  - id: jagged_to_padded_dense_forward
+    description: |
+      Converts a jagged (variable-length) tensor to a padded dense tensor.
+    for:
+      - _jagged_to_padded_dense_forward
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - Tensor
+    stages:
+      - alpha: '5.4'
   - id: kron
     description: Computes the Kronecker product of `input` and `other`.
     for:

src/flag_gems/__init__.py

@@ -61,6 +61,7 @@ def torch_ge(v):
     ("_grouped_mm", group_mm),
     ("_index_put_impl_", _index_put_impl_),
     ("_is_all_true", _is_all_true),
+    ("_jagged_to_padded_dense_forward", _jagged_to_padded_dense_forward),
     ("_linalg_eigvals", _linalg_eigvals),
     ("_log_softmax", log_softmax),
     ("_log_softmax.out", log_softmax_out),

src/flag_gems/ops/__init__.py

@@ -7,6 +7,9 @@
     _functional_sym_constrain_range_for_size,
 )
 from flag_gems.ops._is_all_true import _is_all_true
+from flag_gems.ops._jagged_to_padded_dense_forward import (
+    _jagged_to_padded_dense_forward,
+)
 from flag_gems.ops._linalg_eigvals import _linalg_eigvals
 from flag_gems.ops._resize_output import _resize_output
 from flag_gems.ops._safe_softmax import _safe_softmax
@@ -493,6 +496,7 @@
     "_functional_sym_constrain_range_for_size",
     "_index_put_impl_",
     "_is_all_true",
+    "_jagged_to_padded_dense_forward",
     "_linalg_eigvals",
     "_resize_output",
     "_safe_softmax",

src/flag_gems/ops/_jagged_to_padded_dense_forward.py

@@ -0,0 +1,131 @@
+# Copyright 2026, The 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.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import logging
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.utils import libentry
+from flag_gems.utils import triton_lang_extension as tle
+
+logger = logging.getLogger(__name__)
+
+
+@libentry()
+@triton.jit
+def _jagged_to_padded_dense_forward_kernel(
+    values,
+    offsets,
+    output,
+    padding_value: tl.constexpr,
+    batch_size: tl.constexpr,
+    max_length: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Kernel for converting jagged tensor to padded dense tensor.
+
+    Args:
+        values: 1D tensor containing concatenated variable-length sequences
+        offsets: 1D tensor of start positions for each sequence
+        output: 2D output tensor of shape (batch_size, max_length)
+        padding_value: scalar value for padding
+        batch_size: number of sequences
+        max_length: maximum length of each sequence
+    """
+    pid = tle.program_id(axis=0)
+    batch_idx = pid
+
+    if batch_idx >= batch_size:
+        return
+
+    # Get the start and end offset for this sequence
+    seq_start = tl.load(offsets + batch_idx)
+    seq_end = tl.load(offsets + batch_idx + 1)
+
+    # Calculate the actual sequence length
+    seq_length = seq_end - seq_start
+
+    # Compute the row offset in the output
+    row_offset = batch_idx * max_length
+
+    # Fill with padding value (vectorized per block)
+    for j in tl.range(0, max_length, BLOCK_SIZE):
+        out_offsets = row_offset + j + tl.arange(0, BLOCK_SIZE)
+        out_mask = (j + tl.arange(0, BLOCK_SIZE)) < max_length
+        tl.store(output + out_offsets, padding_value, mask=out_mask)
+
+    # Copy actual values (vectorized per block)
+    for j in tl.range(0, seq_length, BLOCK_SIZE):
+        offsets_vec = seq_start + j + tl.arange(0, BLOCK_SIZE)
+        mask = offsets_vec < seq_end
+
+        values_vec = tl.load(values + offsets_vec, mask=mask, other=padding_value)
+
+        out_offsets = row_offset + j + tl.arange(0, BLOCK_SIZE)
+        out_mask = (j + tl.arange(0, BLOCK_SIZE)) < seq_length
+
+        tl.store(output + out_offsets, values_vec, mask=out_mask)
+
+
+def _jagged_to_padded_dense_forward(values, offsets, max_lengths, padding_value=0.0):
+    """Convert a jagged (variable-length) tensor to a padded dense tensor.
+
+    Args:
+        values: 1D tensor containing concatenated variable-length sequences
+        offsets: List of 1D tensors containing start positions for each sequence
+        max_lengths: List of integers specifying maximum length for each batch dimension
+        padding_value: Value to use for padding (default: 0.0)
+
+    Returns:
+        Padded dense tensor
+    """
+    logger.debug("GEMS JAGGED TO PADDED DENSE FORWARD")
+
+    # Currently only supports single batch dimension
+    if not isinstance(offsets, (list, tuple)):
+        offsets = [offsets]
+    if not isinstance(max_lengths, (list, tuple)):
+        max_lengths = [max_lengths]
+
+    num_batch_dims = len(offsets)
+    assert (
+        num_batch_dims == 1
+    ), f"Only single batch dimension is supported, got {num_batch_dims}"
+
+    # Single batch dimension: 1D values, 1D offsets
+    offsets_0 = offsets[0]
+    batch_size = offsets_0.numel() - 1
+    max_length = max_lengths[0]
+
+    # Compute output shape
+    # For single batch dim: (batch_size, max_length)
+    output_shape = (batch_size, max_length)
+    output = torch.empty(output_shape, dtype=values.dtype, device=values.device)
+
+    grid = lambda meta: (batch_size,)
+    _jagged_to_padded_dense_forward_kernel[grid](
+        values,
+        offsets_0,
+        output,
+        padding_value,
+        batch_size,
+        max_length,
+        # BLOCK_SIZE=128 balances occupancy and memory efficiency for typical sequence lengths
+        BLOCK_SIZE=128,
+    )
+
+    return output

tests/test_jagged_to_padded_dense_forward.py

@@ -0,0 +1,72 @@
+import numpy as np
+import pytest
+import torch
+
+import flag_gems
+
+from . import accuracy_utils as utils
+
+
+@pytest.mark.jagged_to_padded_dense_forward
+@pytest.mark.parametrize("batch_size", [1, 8, 32, 128])
+@pytest.mark.parametrize("max_length", [8, 16, 32, 128])
+@pytest.mark.parametrize("dtype", utils.FLOAT_DTYPES)
+def test_jagged_to_padded_dense_forward(batch_size, max_length, dtype):
+    # Create variable-length sequences
+    # Generate random sequence lengths
+    np.random.seed(42)
+    seq_lengths = np.random.randint(1, max_length + 1, size=batch_size).tolist()
+
+    # Create offsets tensor (cumulative)
+    offsets = [0] + list(np.cumsum(seq_lengths).astype(int).tolist())
+    offsets = torch.tensor(offsets, device=flag_gems.device, dtype=torch.int64)
+
+    # Create values tensor (concatenated sequences)
+    total_length = sum(seq_lengths)
+    values = torch.randn(total_length, dtype=dtype, device=flag_gems.device)
+
+    # Reference implementation
+    ref_values = utils.to_reference(values)
+    ref_offsets = utils.to_reference(offsets)
+
+    ref_out = torch.ops.aten._jagged_to_padded_dense_forward(
+        ref_values, [ref_offsets], [max_length], 0.0
+    )
+
+    # GEMS implementation
+    with flag_gems.use_gems():
+        res_out = torch.ops.aten._jagged_to_padded_dense_forward(
+            values, [offsets], [max_length], 0.0
+        )
+
+    utils.gems_assert_close(res_out, ref_out, dtype)
+
+
+@pytest.mark.jagged_to_padded_dense_forward
+@pytest.mark.parametrize("batch_size", [8, 32])
+@pytest.mark.parametrize("max_length", [16, 32])
+@pytest.mark.parametrize("padding_value", [0.0, -1.0, 1.5])
+def test_jagged_to_padded_dense_forward_padding(batch_size, max_length, padding_value):
+    # Test with different padding values
+    np.random.seed(42)
+    seq_lengths = np.random.randint(1, max_length + 1, size=batch_size).tolist()
+
+    offsets = [0] + list(np.cumsum(seq_lengths).astype(int).tolist())
+    offsets = torch.tensor(offsets, device=flag_gems.device, dtype=torch.int64)
+
+    total_length = sum(seq_lengths)
+    values = torch.randn(total_length, dtype=torch.float32, device=flag_gems.device)
+
+    ref_values = utils.to_reference(values)
+    ref_offsets = utils.to_reference(offsets)
+
+    ref_out = torch.ops.aten._jagged_to_padded_dense_forward(
+        ref_values, [ref_offsets], [max_length], padding_value
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.ops.aten._jagged_to_padded_dense_forward(
+            values, [offsets], [max_length], padding_value
+        )
+
+    utils.gems_assert_close(res_out, ref_out, torch.float32)