[vthumbe1503]

Description

Validate grouped split sizes inside the existing device-side multi-output offset scan. GroupedLinear requests 128-row alignment; the generic grouped MLP offset path requests 256-row alignment. The check runs on every CUDA graph replay without a host read or an added kernel launch.

nvte_splits_to_offsets_multi_v2(split_sizes, split_alignment, outputs, ...) exposes the check. The original C API symbol remains and forwards with alignment 1. The PyTorch private binding accepts an optional split_alignment.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)

Changes

• Validate every device-resident split in the prefix-sum kernel, including later 256-element scan chunks. NVTE_DEVICE_ERROR prints and asserts; __trap() keeps the failure effective in Release builds where NDEBUG removes the assert. NVTE_BUILD_DEBUG=1 enables the assert when building, not at runtime.
• Use a bitmask for power-of-two alignments and run the existing scan with 128 threads for up to 128 splits, retaining 256 threads for larger inputs.
• Pass alignment from GroupedLinear and grouped MLP into the multi-output scan.
• Keep the device-trap failure case out of the automated tests because it invalidates the test process's CUDA context. Existing grouped tensor tests pass; this PR has no negative alignment regression test.
• Document the device-side alignment contract.

Offset-scan GPU time on GB200

Measured the one- and four-output scan with 8 and 16 experts, 256 rows per expert, and split_alignment=128 on the PR. The four-output call uses strides [1,1,7168,4096], matching a grouped linear FC1 shape. Inputs are CUDA int64 splits; the first output is int32 and the remaining outputs are int64. Four-output calls use bulk allocation.

Experts	Outputs	Base µs/call	PR µs/call	PR − Base µs
8	1	2.884	2.913	+0.029
16	1	2.883	2.912	+0.028
8	4	3.651	3.648	−0.003
16	4	3.651	3.652	+0.001

GPU times are medians of seven samples, each timing 40 replays of a CUDA graph containing 64 scan calls. Baseline e32e1238 and PR 06228625 used the same GB200 and container, with their respective built extensions.

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete for paths that run the multi-output scan
• I have commented the device assertion and barrier behavior
• I have made corresponding changes to the API documentation
• My changes generate no new warnings
• I have added tests that prove the behavior
• The focused existing tests pass locally

[vthumbe1503]

/te-ci L0

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge, with no new actionable issues found.

Summary

Adds device-side split alignment checks to the existing offset scan. The original C API keeps its behavior, while PyTorch callers can request alignment.

• Grouped offset scans check each split against its requested alignment.

_{Reviews (6) · Last reviewed commit: "Apply suggestion from @timmoon10" · Reviewed by Greptile}

[timmoon10]

Overall LGTM, with minor suggestions.

[timmoon10]

This is fine, but bitmasking seems overkill for a kernel that's unlikely to be compute-bound. The following is a little simpler.

Suggested change

	// Validate before the scan. All threads leave together so no thread waits at a barrier
	// after another thread has trapped on invalid split metadata.
	const int64_t split =
	idx < args.num_splits ? load_split_size(args.split_sizes, args.split_sizes_dtype, idx) : 0;
	if (args.split_alignment > 1) {
	const size_t alignment_mask = args.split_alignment - 1;
	bool misaligned;
	if ((args.split_alignment & alignment_mask) == 0) {
	misaligned = (static_cast<size_t>(split) & alignment_mask) != 0;
	} else {
	misaligned = static_cast<size_t>(split) % args.split_alignment != 0;
	}
	const bool invalid = idx < args.num_splits && (split < 0 || misaligned);
	if (__syncthreads_or(invalid)) {
	if (tid == 0) {
	NVTE_DEVICE_ERROR("Grouped split size is not divisible by split_alignment.");
	__trap(); // NVTE_DEVICE_ERROR's assert is compiled out of release builds.
	}
	return;
	}
	}
	block_scan[tid] = split;
	__syncthreads();
	// Load input from global memory into shared memory
	if (idx < args.num_splits) {
	block_scan[tid] = load_split_size(args.split_sizes, args.split_sizes_dtype, idx);
	} else {
	block_scan[tid] = 0;
	}
	if (args.split_alignment > 1) {
	// Check that splits have required alignment
	if (__syncthreads_or(block_scan[tid] % args.split_alignment)) {
	if (tid == 0) {
	NVTE_DEVICE_ERROR("Grouped split size is not divisible by split_alignment.");
	__trap(); // NVTE_DEVICE_ERROR's assert is compiled out of release builds.
	}
	return;
	}
	} else {
	__syncthreads();
	}

[vthumbe1503]

I evaluated this change. And not doing the bit masking costs us 0.07us consistently for the kernel.

[timmoon10]

Is there reason to support separate kernels for these two cases? This kernel seems launch-bound, so the smaller variant doesn't give us much perf benefit.

[vthumbe1503]

The improves perf by about 0.2 us. I would assume this is purely warp scheduling overhead (4 vs 8 warps). But at the scale of 2-3 us for the kernel that is 10% perf benefit. And so I kept this change. But you can see that we are loosing the perf by about 0.2 us due to the introduction of alginment check.