[Performance] repeat_interleave.Tensor ignores output_size and still reads the output length from device #7241

Describe the bug

The common FlagGems repeat_interleave_tensor accepts output_size but still discovers the output length with cumsum[-1].item(). On the R5 DCU Tensor overload, supplying output_size=854524 gives 272.25 us, versus 270.69 us without it; the all-cuda control drops to 20.36 us with known size. The values API also pays a separate Torch-FL index-check cost. All four cases are correct.

Environment

• Reproduced on 2026-10-18, Linux, one Hygon BW / gfx936 DCU, 80 CUs, approximately 64 GiB; all A/B/C measurements used the same physical device within Slurm job 936955, completed with exit status 0.
• Python 3.10.20; DTK 26.04; HIP 6.3.26113.
• Shared deployment: r5-20260928-tfl6d84d12b-fge7b4a865.
• Torch-FL 2.10.0+das.dtk2604.glibc228.g6d84d12b, commit 6d84d12b715c307e5b660fc33e0f51e660602f6a.
• FlagGems 5.4.0rc2.post1+ge7b4a865f, commit e7b4a865fce6d85861ee91a6aca56564ef9acf7d.
• FlagTree / Triton implementation: 0.6.0+hcu3.6.glibc228.compat1.git8a1dda94, HCU target; no separate triton distribution version was reported by package metadata.
• B/C: official PyTorch 2.10.0+cpu, git 449b1768410104d3ed79d3bcfe4ba1d65c7f22c0, with R5's DCU HIP bundle and compatibility layer.
• A: native DAS PyTorch 2.10.0+das.opt1.dtk2604.2606151054.g98b1a8, git 98b1a89bbf5bcf9572ae9e23279e6ecee92857ed; cuda:0; Torch-FL was not imported.
• Build provenance: the published R5 prebuilt wheels were used, with bundle_mode=official-core-decoupled and local_patches=[]; no install, rebuild, or shared Release modification was performed for this experiment. The original wheel compiler command is not reconstructed here.

Comparability limitation: A and B/C do not use the same Torch frontend binary. Their HIP libraries have matching ELF Build IDs but different RUNPATH/file hashes. Therefore A/B is an end-user stack comparison, not a single-variable kernel comparison. B/C uses the same Torch frontend and R5 binaries and isolates route selection, except for hand-written safety/boxing wrappers which remain active.

Reproduction

The following two-file reproducer is the unchanged hardware-tested demo and timer. Save both files in the same directory. It does not depend on the private cluster paths or the rest of the benchmark suite. Import order in common.py is torch_fl before torch for --device flagos:0.

05_repeat_interleave.py

"""Expand 17096 rows to 854524 items, with/without known output size."""
from common import setup, check, run

torch, device, args = setup(__doc__, {'rows': 17096, 'items': 854524})
r_cpu = torch.full((args.rows,), args.items // args.rows, dtype=torch.int64)
r_cpu[:args.items % args.rows] += 1
x_cpu = torch.arange(args.rows)
x, repeats = x_cpu.to(device), r_cpu.to(device)
expected = torch.repeat_interleave(x_cpu, r_cpu)
cases = {
    'repeat_values_dynamic': lambda: torch.repeat_interleave(x, repeats),
    'repeat_values_known': lambda: torch.repeat_interleave(x, repeats, output_size=args.items),
    'repeat_Tensor_dynamic': lambda: torch.repeat_interleave(repeats),
    'repeat_Tensor_known': lambda: torch.repeat_interleave(repeats, output_size=args.items),
}
for fn in cases.values():
    check(torch, fn(), expected)
run(torch, device, args, cases)

"""Shared CLI/timer only. No sparse/scatter dependency; attach with one demo."""
import argparse
import importlib.metadata
import importlib.util
import json
import os
import re
import socket
import statistics
import sys
import time


def positive(value):
    value = int(value)
    if value < 1:
        raise argparse.ArgumentTypeError('must be positive')
    return value


def setup(description, sizes):
    p = argparse.ArgumentParser(description=description)
    p.add_argument('--device', default='auto', help='auto, flagos[:N], cuda[:N], or cpu')
    for key, default in sizes.items():
        p.add_argument('--' + key.replace('_', '-'), type=positive, default=default)
    p.add_argument('--warmup', type=positive, default=10)
    p.add_argument('--iters', type=positive, default=50)
    p.add_argument('--repeats', type=positive, default=3)
    p.add_argument('--seed', type=int, default=1234)
    p.add_argument('--case', help='time only this case; see the demo source for names')
    p.add_argument('--profile', action='store_true', help='print CPU profiler table of the first selected case')
    a = p.parse_args()
    if not re.fullmatch(r'auto|cpu|(?:flagos|cuda)(?::[0-9]+)?', a.device):
        p.error('--device must be auto, cpu, flagos[:N], or cuda[:N]')
    if a.device.startswith('flagos'):
        import torch_fl  # Must precede torch.
    elif a.device == 'auto' and importlib.util.find_spec('torch_fl') is not None:
        try:
            import torch_fl
        except ImportError:
            pass
    import torch
    def available(name):
        try:
            torch.empty(0, device=name)
            return True
        except (RuntimeError, TypeError, ModuleNotFoundError):
            return False
    if a.device == 'auto':
        a.device = next((d for d in ['flagos:0', 'cuda:0'] if available(d)), 'cpu')
    elif a.device in ['cuda', 'flagos']:
        a.device += ':0'
    if not available(a.device):
        raise RuntimeError(f'{a.device} is unavailable in this environment')
    d = torch.device(a.device)
    if d.type in ['cuda', 'flagos']:
        getattr(torch, d.type).set_device(d.index)
    a.generator = torch.Generator(device='cpu').manual_seed(a.seed)
    def sync():
        if d.type == 'cuda':
            torch.cuda.synchronize(d)
        elif d.type == 'flagos':
            if hasattr(torch.flagos, 'synchronize'):
                torch.flagos.synchronize(d)
            elif hasattr(torch, 'accelerator'):
                torch.accelerator.synchronize()
            else:
                raise RuntimeError('No FlagOS synchronize API; refusing asynchronous timing')
    a.sync = sync
    return torch, d, a


def check(torch, actual, expected):
    # Run correctness checks outside the timed region.
    actual = actual.detach().cpu()
    assert actual.dtype == expected.dtype and actual.shape == expected.shape
    assert torch.equal(actual, expected), 'output differs from CPU reference'


def run(torch, d, a, cases):
    if a.case:
        cases = {k: v for k, v in cases.items() if k == a.case}
        if not cases:
            raise ValueError('Unknown --case; see the demo source for valid names')
    metadata = {'host': socket.gethostname(), 'device': str(d), 'torch': torch.__version__,
                'hip': torch.version.hip, 'python': sys.executable, 'torch_file': torch.__file__,
                'warmup': a.warmup, 'iters': a.iters, 'repeats': a.repeats, 'seed': a.seed,
                'sizes': {k: v for k, v in vars(a).items() if type(v) is int},
                'torch_fl_imported': 'torch_fl' in sys.modules,
                'flag_gems_imported': 'flag_gems' in sys.modules,
                'env': {k: v for k, v in os.environ.items() if k.startswith('FLAGOS_OP_')
                        or k in ['FLAGOS_FORCE_BACKEND', 'USE_FLAGGEMS', 'HIP_VISIBLE_DEVICES',
                                 'CUDA_VISIBLE_DEVICES', 'ROCR_VISIBLE_DEVICES', 'SLURM_JOB_ID']}}
    for package in ['torch', 'torch-fl', 'flag-gems', 'flagtree']:
        try:
            metadata[package + '_distribution'] = importlib.metadata.version(package)
        except importlib.metadata.PackageNotFoundError:
            pass
    print(json.dumps(metadata, ensure_ascii=False), flush=True)
    for name, fn in cases.items():
        for _ in range(a.warmup):
            fn()
        samples = []
        for _ in range(a.repeats):
            a.sync()
            start = time.perf_counter()
            for _ in range(a.iters):
                fn()
            a.sync()
            samples.append((time.perf_counter() - start) * 1e6 / a.iters)
        rounds = ', '.join(f'{value:.2f}' for value in samples)
        print(f'{name:36} {statistics.median(samples):10.2f} us  rounds_us=[{rounds}]', flush=True)
    if a.profile:
        name, fn = next(iter(cases.items()))
        with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as prof:
            fn()
            a.sync()
        print(f'\nPROFILE {name}', flush=True)
        print(prof.key_averages().table(sort_by='self_cpu_time_total', row_limit=50), flush=True)

# In a clean, activated R5 environment. These do not install or modify packages.
export TORCH_DEVICE_BACKEND_AUTOLOAD=0
export HIP_VISIBLE_DEVICES=0 CUDA_VISIBLE_DEVICES=0
export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1
unset FLAGOS_USE_FLAGGEMS FLAGOS_USE_FLAGGEMS_CPP USE_FLAGGEMS
unset FLAGOS_FORCE_BACKEND FLAGOS_BACKEND_CONFIG
for key in $(compgen -e | grep '^FLAGOS_OP_' || true); do unset "$key"; done
# R5 activation selects: GEMS_VENDOR=hygon, FLAGTREE_BACKEND=hcu,
# FLAGOS_USE_FLAGTREE=1. Retain its correct DTK/HIP loader configuration.

FLAGOS_USE_FLAGGEMS was unset; it was not used as a routing switch. B used R5's shipped per-op configuration. C used a private copy of that configuration with every route entry changed to cuda, not FLAGOS_FORCE_BACKEND=vendor.

Generate the C configuration once in the activated R5 environment. This only writes a newly created temporary directory, never the installed configuration:

BOXING_CONF=$(python - <<'PY'
import importlib.util, pathlib, re, tempfile
spec = importlib.util.find_spec("torch_fl")
source = pathlib.Path(next(iter(spec.submodule_search_locations))) / "configs/backends_dcu.conf"
text = source.read_text()
pattern = r"^(\s*[A-Za-z0-9_.]+\s*=\s*)[A-Za-z0-9_]+"
entries = re.findall(pattern, text, flags=re.M)
assert entries
text = re.sub(pattern, r"\1cuda", text, flags=re.M)
assert len(re.findall(r"^[A-Za-z0-9_.]+\s*=\s*cuda\s*(?:#.*)?$", text, re.M)) == len(entries)
target = pathlib.Path(tempfile.mkdtemp(prefix="fgperf-")) / "boxing.conf"
target.write_text(text)
print(target)
PY
)

# B: default R5. Run in a fresh process.
python 05_repeat_interleave.py --device flagos:0 --warmup 50 --iters 200 --repeats 7
# C: same R5 binaries, all-cuda table; special C++ wrappers still execute.
FLAGOS_BACKEND_CONFIG="$BOXING_CONF" python 05_repeat_interleave.py --device flagos:0 --warmup 50 --iters 200 --repeats 7
# A: run in a SEPARATE native DAS environment with no Torch-FL autoload.
TORCH_DEVICE_BACKEND_AUTOLOAD=0 python 05_repeat_interleave.py --device cuda:0 --warmup 50 --iters 200 --repeats 7

The reported table was collected with the full original run_all.py suite in each group process, not one new process per individual case. The standalone demo above preserves the tested source/workload, but exact absolute timings may depend on preceding suite cases. The umbrella report supplies the complete suite for protocol-identical reruns. Repeat in three process rounds in the orders stated above. The --case option can select one case. To verify routes, use a separate diagnostic process with FLAGOS_LOG=dispatch,fallback and short warmup/iters/repeats; never use that output as a formal latency sample. When testing another release, verify its actual routes before comparing results.

Expected behavior

The implementation should preserve PyTorch correctness while avoiding the identified unnecessary algorithm/known-shape overhead; a default accelerator route should not silently impose this magnitude of steady-state regression for the reported case. Exact performance acceptance is to be agreed with maintainers.

Actual behavior and measured output

For the values API with known size, a separate device diagnostic recorded A 4 kernels / 34.72 us, B 7 / 105.28 us, C 5 / 49.76 us. All four cases matched CPU exactly. No exception occurred; the benchmark job exited 0.

Measurement: eager execution, not torch.compile. Inputs were created on CPU, seed 1234, and moved to the device before correctness checks, warmup, and timing. Original parameters were warmup=10 / iters=50 / repeats=3. The table below uses warmup=50 / iters=200 / repeats=7, three independent processes per group, ordered A/B/C, C/A/B, B/C/A. Reported values are the median of the three process medians, in microseconds per API call. Each measured batch has explicit device synchronization before/after it. Correctness checks and H2D transfers are outside timing; synchronization intrinsically performed inside an API is included.

Formal timing did not enable profiler, dispatch logging, or Python tracing. Warmed diagnostic compile/autotune auditing recorded no corresponding calls in the audited invocations. Device timelines were collected separately: A used Kineto, B/C used HIP profiler with warmed HIPTX ranges. These are different instrumentation methods; kernel counts and device-time orders of magnitude support the mechanism, not an exact wall-clock-minus-device decomposition. No cross-operator graph fusion was active.

Root-cause evidence and ownership

Observed chain: deterministic CPU int64 repeats (17096,), sum 854524, values arange(17096) -> pre-timing copy -> aten::repeat_interleave.Tensor(Tensor repeats, *, SymInt? output_size=None) on PrivateUse1 -> Torch-FL generated Python bridge -> common FlagGems repeat_interleave_tensor; the values API is Composite and then calls aten::index_select.

Confirmed source issue: repeat_interleave_tensor accepts output_size but never uses it. It unconditionally runs repeats.cumsum(axis=0) and result_size = cumsum[-1].item() before allocating its result. The prefix sum may still be required to construct offsets; the avoidable part reported here is determining the known output length via a device scalar readback. Do not infer that the entire prefix sum can simply be removed.

The deployed generated C++ bridge does forward output_size in kwargs; this is not a Torch-FL argument-loss report. Generated bridge. The values API's index_select additionally pays Torch-FL's index range-check cost, tracked separately as flagos-ai/Torch-FL#579. Even all-cuda C retains this wrapper, explaining why C values-known is not as fast as A; that overhead must not be charged again to this FlagGems fix.

Ownership: FlagGems common repeat_interleave implementation for known-size handling. The broader small-op wrapper cost is related to #6570 and #4812, but the total latency is not attributed solely to output_size.

Suggested implementation direction

Use a supplied valid output_size to size the allocation without host-reading the cumulative-sum tail. Preserve the prefix offsets needed to generate the output, and preserve required validation for negative repeats and inconsistent sizes. Ensure both Tensor and values APIs forward/benefit from the known size. Avoid a new synchronous host validation that negates the optimization. Do not include index-wrapper savings as repeat_interleave-only savings.

Verification plan

• Unit: Tensor and values overloads, with/without valid output_size; zero/negative repeats and invalid output_size should follow the operator contract.
• Integration: verify the actual output_size reaches the FlagGems implementation and no tail-scalar readback is used only to discover an already known shape.
• Manual DCU: repeat the four-case A/B/C matrix on one physical device; profile in a separate process; report prefix-scan, known-length readback and values/index-selection effects separately. Correctness must remain exact and known-size steady-state latency should improve without attributing the separate index-check optimization to this fix.

Context and current-source check

Read-only upstream inspection on 2026-10-18: FlagGems master f89952b8afd1fdbf44e99265ec7835c3aa2bb532; Torch-FL main 3edbbf91bc9d2411b963ca9111bc018bc7769adc. The relevant Hygon sort, common repeat_interleave, and Torch-FL index_bounds files are byte-identical to the deployed snapshots used here. This supports source applicability, not a new hardware benchmark of current upstream. No performance fix for this reported case has been established as released/deployed/verified by this investigation.

Both open/closed Issues and PRs were searched in both repositories using the operator names, Hygon, CheckIndex/index_bounds/bounds, performance, and related issue references. Plausible matches were inspected, including FlagGems #6921 (MetaX sort), #2974 (repeat correctness), #6570/#4812 (shared overhead), backend-specific PRs #6384/#6507, and Torch-FL #381/#395 (index safety), #252 (add route mitigation). These reports are related context, not proof that the current cases are fixed. The canonical team Feishu tracker could not be read; the fallback snapshot is dated 2026-09-23, so team-tracker duplicate coverage is limited.

Prepared by Codex at the user's request from an existing hardware reproduction; exact model identifier is not exposed by the session.