[FEA]: CTA-level load-balanced segmented scan for irregularly-sized segments

[kierandidi]

Is this a duplicate?

• I confirmed there appear to be no duplicate issues for this request and that I agree to the Code of Conduct

Area

General CCCL

Is your feature request related to a problem? Please describe.

This feature request relates to tmol, a GPU-accelerated biomolecular energy function in PyTorch with custom C++/CUDA kernels. The CUDA code there currently depends on moderngpu, which is no longer actively maintained.

The most critical moderngpu dependency is a load-balanced segmented scan used for batched forward/inverse kinematics on protein backbone chains. Proteins are modeled as kinematic trees, and computing atom positions requires a segmented prefix scan with 4x4 matrix composition as the operator, where each segment is a protein chain of different length.

The problem is irregular parallelism: a batch of proteins contains segments of wildly different lengths (10 atoms for a small peptide up to 1000+ for a large protein, with 50-200 segments per batch). Naive approaches (one CTA per segment) are extremely load-imbalanced. We need CTA-level building blocks for load-balanced segmented scans, which moderngpu provides but CCCL currently does not.

Describe the solution you'd like

A CCCL-native API for load-balanced segmented scan that supports:

1. Merge-path based load balancing across CTAs — data elements AND segment boundaries are interleaved into a merged sequence so each CTA gets roughly nt * vt work items regardless of segment sizes (equivalent to moderngpu's cta_load_balance_t).

2. CTA-level segmented scan — within each CTA, threads process their vt work items sequentially, resetting the accumulator at segment boundaries detected via merge flags. A CTA-wide segmented scan then combines results across threads while respecting segment boundaries (equivalent to moderngpu's cta_segscan_t).

3. Multi-CTA carry-out propagation (spine scan) — carry-out values from all CTAs are scanned recursively (still segment-aware), with a downward sweep propagating accumulated carry-in values back to each CTA's elements, stopping at the first segment boundary in each CTA.

4. User-defined non-commutative types and operators — our scan operates on 4x4 homogeneous transformation matrices (16 floats) with matrix composition as the operator. This is associative but not commutative.

5. Segment boundaries as sorted offset arrays (not boolean head flags).

6. Both inclusive and exclusive scan modes.

The current implementation is ~400 lines built on moderngpu's CTA-level primitives: kernel_segscan.cuh. It uses:

• cta_load_balance_t for merge-path work distribution (lines 296-308)
• cta_segscan_t for CTA-level segmented scan (lines 79-86)
• cta_reduce_t for finding the first segment boundary in a CTA (lines 74-75)
• Recursive spine_segreduce for multi-CTA propagation (lines 132-217)

The call site for the kinematic tree scan is in compiled.cuda.cu, where it scans over 4x4 homogeneous transformation matrices with a custom composition operator.

Ideally, the CCCL equivalent would look something like:

// Hypothetical API
cub::DeviceSegmentedScan::InclusiveScan(
d_temp_storage, temp_storage_bytes,
d_in, d_out,
num_items, num_segments, d_segment_offsets,
compose_op, identity,
stream,
/* load_balanced = */ true // use merge-path load balancing
);

Or as composable CTA-level building blocks similar to moderngpu's approach, which would allow custom per-element functors during the scan (our implementation uses a user-supplied functor f(index, seg, rank) to generate input values on-the-fly rather than reading from a flat array).

Describe alternatives you've considered

• cub::DeviceSegmentedReduce: Only supports reduction, not scan (no prefix sums). Also doesn't do load balancing across segments.
• cub::DeviceScan + manual segmentation: Requires pre-/post-processing to handle segment boundaries, losing the efficiency of the fused load-balanced approach.
• cub::BlockScan with manual decomposition: Possible but would require reimplementing the merge-path load balancing, carry-out propagation, and spine scan — essentially rebuilding what moderngpu already provides.
• Keep vendoring moderngpu: Works but moderngpu is unmaintained since 2021 and doesn't support newer architectures natively.

Additional context

• The data type being scanned is a 16-float struct (4x4 matrix) with a non-commutative composition operator — not just int with plus.
• Typical workload: O(100K) total elements across O(100) irregularly-sized segments.
• The full codebase using this primitive is publicly available at github.com/uw-ipd/tmol, specifically:
  • Segmented scan implementation: kernel_segscan.cuh
  • Usage in kinematic tree traversal: compiled.cuda.cu
  • Device-level scan/reduce wrappers: device_operations.cuda.impl.cuh

[github-actions]

Hi @kierandidi!

Thanks for submitting this issue - the CCCL team has been notified and we'll get back to you as soon as we can!
In the mean time, feel free to add any relevant information to this issue.

[gevtushenko]

@kierandidi to make sure I understand correctly: what you ultimately need is an efficient cub::DeviceSegmentedScan::InclusiveScan, rather than CTA-level primitives for load balancing. Is that right?

@oleksandr-pavlyk recently added cub::DeviceSegmentedScan::InclusiveSegmentedScan, and there's also an optimization #6712 for segmented scan currently in progress.

If you could share a standalone benchmark demonstrating how load imbalance in the current implementation leads to a significant slowdown compared to your algorithm, that would be extremely valuable. It would help us evaluate and choose among the optimization strategies we're considering for that PR.

Finally, if we're able to get segmented scan to SOL performance, would you still have a need for the CTA-level load-balancing primitives?

[kierandidi]

@gevtushenko thanks for the reply! Regarding your questions:
1 and 3: yes, if there was a segmented scan to SOL performance that fits the API we describe then we would not need the CTA-level load-balancing primitives.
2: I do not have the capacity right now to implement the current Cub inclusive segmented scan for our application, but I could provide a small benchmark for our current implementation that if someone from your end tries to implement with the current cub primitives should show this problem since the inputs are very inbalanced if that is helpful?

[fbusato]

@gevtushenko the strategy that @kierandidi is proposing makes a lot of sense. The current strategy in SegmentedScan involves mapping each thread block to exactly one segment. This is clearly suboptimal for any workload where segments are not similar to the block size.

The Merge-path based load balancing is almost optimal considering that the operator (compose_op) is computationally inexpensive. The main drawback is that with merge-path you need to pay the price for the partition phase (additional kernel).
Also, merge-path could generate empty work when there are zero segments.

An even better strategy could be assign the same number of element per thread and generate the segment "flags" implicitly with a binary search.

[gevtushenko]

1 and 3: yes, if there was a segmented scan to SOL performance that fits the API we describe then we would not need the CTA-level load-balancing primitives.

Got it, @kierandidi thank you for the clarification! Apart from the load balancing, there are many optimizations that we can do on segmented scan end. Long-term, it's better for users to rely on it instead of building generic device-level primitives from block-level components. I'm making this a sub-issue on a larger segmented-scan optimization theme #7850.

I could provide a small benchmark for our current implementation that if someone from your end tries to implement with the current cub primitives should show this problem since the inputs are very inbalanced if that is helpful?

That'd be immensely helpful. Please, provide details on #7852.

The current strategy in SegmentedScan involves mapping each thread block to exactly one segment.

@fbusato we won't assign a single thread block per segment after #6712, but I agree, we should investigate merge path so that users don't have to build this algorithm from block-level primitives if their use case suffers from load-balancing issues.

[seanbaxter]

It's worthwhile to provide device- and block-level load-balancing search compared to just doing segmented scan. The load balancing search handles diverse problems related to sparsity, like SpM*V and database inner/outer joins. I might dust this stuff off and make a PR.

[fbusato]

Thanks @seanbaxter, I fully support your proposal. Your ModernGPU library, which leverages merge-path technology, has been an invaluable resource for the community. For SpM*V kind of problems, merge-path is not the optimal solution, but I can help here with my background on sparse linear algebra.