Skip to content

[AI][BUG][DCU] Track FlagGems scatter_add_ failure with stride-zero expanded index #349

Description

@wzx226226

Issue Type

  • Bug Report (Torch-FL / FlagGems operator integration)

AI Agent Information

  • Agent: Codex.
  • Model: runtime model identifier was not recorded in the original experiment artifacts.
  • Session context: this issue records a DCU mixed-FlagGems correctness failure first observed in DiffDock. The Tensor-only reproduction was measured on 2026-09-18; it is not a claim of a new hardware test against an unmodified current upstream HEAD.

Summary

On Hygon DCU, aten::scatter_add_.default returns incorrect values when its index is a valid PyTorch view(...).expand(...) Tensor with stride (1, 0). DAS boxing produces the correct PyTorch result for the same input, while FlagGems mismatches all 120 output elements in the deterministic reproducer.

The normal unpadded expand layout only owns 340 index elements although its logical shape is (340, 12). The affected FlagGems fast path treats that index as contiguous and reads addresses through logical offset 4079. Depending on surrounding device allocations, this can either read arbitrary valid memory and silently compute a wrong result, or cross an inaccessible device page and abort with a VMFault.

Ownership: Torch-FL correctly routes aten::scatter_add_.default to FlagGems. The error is in FlagGems' 2-D scatter_add_ addressing, which ignores index strides. This Torch-FL tracker records the model-facing DCU failure and acceptance criteria for a FlagGems repair; it does not attribute the kernel addressing defect to DiffDock.

Environment (for bug reports)

  • Hardware/platform: Hygon DCU / BW1000, physical device 0.
  • Runtime environment: FlagOS DCU environment provided with the tested release.
  • PyTorch runtime: 2.10.0 with Torch-FL DCU integration.
  • Torch-FL mixed configuration: backends_dcu_flaggems.conf.
  • FlagGems vendor route: hygon, flagos_python.
  • Device: flagos:0.
  • Hardware reproduction date: 2026-09-18.

The report is evidence for the measured release, not a claim that the current upstream main branch was rebuilt and re-tested.

Reproduction

Save as repro_scatter_add_expanded_index.py and run each route in a fresh Python process:

python repro_scatter_add_expanded_index.py boxing
python repro_scatter_add_expanded_index.py flaggems
import os
import sys

route = sys.argv[1]
assert route in {"boxing", "flaggems"}

# Select the route before importing torch_fl.
for key in list(os.environ):
    if key.startswith("FLAGOS_OP_"):
        os.environ.pop(key)
os.environ.pop("FLAGOS_BACKEND_CONFIG", None)
os.environ["HIP_VISIBLE_DEVICES"] = "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["FLAGOS_USE_FLAGGEMS"] = "1" if route == "flaggems" else "0"
os.environ.pop("FLAGOS_USE_FLAGGEMS_CPP", None)
if route == "boxing":
    os.environ["FLAGOS_OP_scatter_add_"] = "cuda"

import torch
import torch_fl

rows, columns, output_rows = 340, 12, 10
device = "flagos:0"
output = torch.zeros((output_rows, columns), dtype=torch.float32, device=device)
source = torch.ones((rows, columns), dtype=torch.float32, device=device)

# This index is logically (340, 12), physically a 340-element expanded view,
# and has stride=(1, 0). The controlled tail makes stride-ignorant reads
# deterministic without requiring a VMFault to reproduce the functional bug.
storage = torch.zeros(rows * columns, dtype=torch.int64, device=device)
storage[:rows] = torch.arange(rows, dtype=torch.int64, device=device) % output_rows
index = storage[:rows].view(rows, 1).expand(rows, columns)
assert index.shape == (rows, columns)
assert index.stride() == (1, 0)

actual = output.scatter_add_(0, index, source)
torch.flagos.synchronize()
expected = torch.zeros((output_rows, columns), dtype=torch.float32)
expected.scatter_add_(0, index.cpu(), source.cpu())
actual_cpu = actual.cpu()
different = actual_cpu != expected

print("route:", route)
print("actual[0, :]  :", actual_cpu[0].tolist())
print("expected[0, :]:", expected[0].tolist())
print("mismatched:", different.sum().item(), "/", different.numel())
torch.testing.assert_close(actual_cpu, expected)

The reproducer is self-contained. The backing storage is deliberately oversized so the incorrect contiguous-address calculation reads controlled data and produces a deterministic wrong answer. It validates the semantic bug without depending on device-page protection behavior.

Expected vs Actual Behavior

Route Result for output row 0 Comparison with CPU reference
DAS boxing [34.0] * 12 Passes; 0 mismatches out of 120
FlagGems [317.0, 311.0, 316.0, ...] Fails; 120 mismatches out of 120

Measured boxing output:

actual[0, :]  : [34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0]
expected[0, :]: [34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0]
mismatched: 0 / 120

Measured FlagGems output:

actual[0, :]  : [317.0, 311.0, 316.0, 311.0, 317.0, 312.0, 318.0, 312.0, 318.0, 312.0, 318.0, 312.0]
expected[0, :]: [34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0, 34.0]
mismatched: 120 / 120
AssertionError: Tensor-likes are not close!
Greatest absolute difference: 284.0 at index (0, 6)

In the original DiffDock input using ordinary unpadded index.view(rows, 1).expand(rows, columns) storage, the same failure was also observed as:

>>>>>>>> KERNEL VMFault !!!! <<<<<<
HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION
SIGABRT / core dump

The VMFault is an additional manifestation of the same invalid address calculation, not the only success criterion. A silently wrong result is already a correctness failure.

Root Cause Analysis / Bottom-layer Execution

The affected FlagGems implementation is scatter_add_2d_kernel in src/flag_gems/ops/scatter_add.py.

The kernel reconstructs an index address as if index had contiguous stride (12, 1):

row = offsets // idx_ncols
col = offsets % idx_ncols
idx_offsets = row * idx_ncols + col
idx = tl.load(index_ptr + idx_offsets, mask=mask, other=0)

For the valid expanded Tensor in this report, the correct addressing is:

index_ptr + row * index_stride0 + col * index_stride1
index_ptr + row * 1 + col * 0

The current formula instead uses row * 12 + col. At the final logical element it tries to read index offset 339 * 12 + 11 = 4079. An ordinary expanded index has only 340 physical int64 elements. Values read through those invalid offsets are then treated as destination rows for tl.atomic_add, which explains both corrupt results and possible device-memory aperture violations.

The execution chain is:

output.scatter_add_(0, index, source)
  -> aten::scatter_add_.default
  -> Torch-FL mixed FlagGems route
  -> FlagGems scatter_add_ wrapper
  -> scatter_add_2d_kernel
  -> contiguous index addressing ignores stride=(1, 0)
  -> invalid index read -> wrong atomic destination or VMFault

Related Code Locations

  • Torch-FL route configuration: torch_fl/configs/backends_dcu_flaggems.conf, scatter_add_ route.
  • FlagGems operation: src/flag_gems/ops/scatter_add.py.
  • Affected kernel: scatter_add_2d_kernel and the idx_offsets address calculation.

Proposed Solution / Upstream Regression Plan

  1. Pass index_stride0 and index_stride1 to the 2-D kernel and load the index using stride-aware addressing:

    idx_offset = row * index_stride0 + col * index_stride1
    idx = tl.load(index_ptr + idx_offset, mask=mask, other=0)
  2. Alternatively, guard the contiguous-only fast path with an explicit layout check and send non-contiguous/expanded index tensors to a tested stride-aware implementation.

  3. Preserve ATen behavior for valid views; do not reject an ordinary expand merely because it has stride zero.

  4. Add regression cases for contiguous, transpose/non-contiguous, stride-zero expanded, and empty indices. Validate against CPU and DAS boxing.

  5. Run a memory-safety profiler or hardware fault test with unpadded expanded storage after functional correctness passes, because silent incorrect writes are as important as visible VMFaults.

Torch-FL Integration Approach

Until a native FlagGems repair is integrated, Torch-FL may use DAS boxing as a temporary safety route for scatter_add_. This must be classified as a forced safety fallback, not as native FlagGems support. The acceptance criterion is that the stride-zero expanded-index reproducer passes on the FlagGems route with no FLAGOS_OP_scatter_add_=cuda override.

Verification Plan

  • Run this exact deterministic stride (1, 0) reproducer on FlagGems without a forced boxing override.
  • Compare all 120 output elements exactly with CPU and DAS boxing.
  • Run the ordinary unpadded expand layout under a memory-safety/device-fault profiler after functional tests pass.
  • Test other dimensions, scatter axes, dtypes, duplicate indices, non-contiguous source/output views, and empty tensors.
  • Re-run the affected DiffDock workload and compare model output with boxing before reporting the integration as fixed.

Checklist

  • DCU hardware environment and reproduction date recorded.
  • Self-contained deterministic reproducer included.
  • Exact stride-zero expanded-index layout is asserted.
  • DAS boxing control and strict full-output comparison included.
  • Silent wrong-result and VMFault manifestations clearly separated.
  • Root cause isolates an index-stride addressing error.
  • No private output paths are required to reproduce or assess the report.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions