fix(scatter_add): accumulate duplicate destination indices via unique_indices flag

src/nkilib_src/nkilib/experimental/misc/scatter_add.py

@@ -32,7 +32,7 @@
 
 
 @nki.jit
-def scatter_add(input: nl.ndarray, dim: int, index: nl.ndarray, src: nl.ndarray) -> nl.ndarray:
+def scatter_add(input: nl.ndarray, dim: int, index: nl.ndarray, src: nl.ndarray, unique_indices: bool = True) -> nl.ndarray:
     """
     Scatter-add from src into input based on indices using gather-accumulate-scatter pattern.
 
@@ -51,6 +51,10 @@ def scatter_add(input: nl.ndarray, dim: int, index: nl.ndarray, src: nl.ndarray)
         dim (int): Dimension along which to scatter (must be 0)
         index (nl.ndarray): [K], 1D tensor of row indices into input
         src (nl.ndarray): [K, D], Source values to scatter-add
+        unique_indices (bool): If True (default), assume destination indices are unique
+            within each 128-row tile and use the fast tile-wide gather/scatter path.
+            If False, process one index per tile so repeated indices accumulate
+            correctly (slower; needed for embedding-gradient-style scatters).
 
     Returns:
         input (nl.ndarray): [N, D], The input tensor with scattered values added
@@ -59,7 +63,8 @@ def scatter_add(input: nl.ndarray, dim: int, index: nl.ndarray, src: nl.ndarray)
         - Input and src tensors must be 2D
         - Index tensor must be 1D
         - dim must be 0
-        - Indices within a tile of 128 rows should be unique for correctness
+        - With unique_indices=True (default) indices within a 128-row tile must be unique;
+          duplicates are silently dropped. Pass unique_indices=False when they may repeat.
 
     Pseudocode:
         for k_tile in tiles(K):
@@ -76,7 +81,10 @@ def scatter_add(input: nl.ndarray, dim: int, index: nl.ndarray, src: nl.ndarray)
 
     _validate_scatter_add_inputs(input, dim, index, src, num_shards)
 
-    k_tile_size = nl.tile_size.pmax
+    # When indices may repeat within a 128-row tile, a tile-wide gather/scatter is
+    # last-write-wins and silently drops duplicate contributions. k_tile_size=1 makes
+    # each index its own sequential tile, so nl.sequential_range accumulates dups correctly.
+    k_tile_size = nl.tile_size.pmax if unique_indices else 1
     d_tile_size = nl.tile_size.psum_fmax * _D_TILE_FACTOR
     num_k_tiles = div_ceil(k_size, k_tile_size)
 

src/nkilib_src/nkilib/experimental/misc/scatter_add_torch.py

@@ -22,6 +22,7 @@ def scatter_add_torch_ref(
     dim: int,
     index: torch.Tensor,
     src: torch.Tensor,
+    unique_indices: bool = True,
 ) -> dict[str, torch.Tensor]:
     """
     PyTorch reference implementation of the scatter_add kernel.

test/integration/nkilib/experimental/misc/test_scatter_add_dup_indices.py

@@ -0,0 +1,84 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+# You may not use this file except in compliance with the License.
+
+"""Regression tests for scatter_add with duplicate indices within a tile.
+
+The base scatter_add test (test_scatter_add.py) constructs indices as per-tile
+permutations, so every 128-row tile has unique destination rows -- it never
+exercises duplicates. Real LM embedding-gradient batches repeat rows constantly
+(common tokens, BOS/EOS), which the tile-wide gather/scatter drops silently.
+These tests pass unique_indices=False and feed sampled-with-replacement indices.
+"""
+
+import nki.language as nl
+import numpy as np
+import pytest
+
+from nkilib_src.nkilib.experimental.misc.scatter_add import scatter_add
+from nkilib_src.nkilib.experimental.misc.scatter_add_torch import scatter_add_torch_ref
+from test.utils.common_dataclasses import CompilerArgs, Platforms
+from test.utils.pytest_parametrize import pytest_parametrize
+from test.utils.pytest_test_metadata import pytest_marks, pytest_test_metadata
+from test.utils.test_orchestrator import Orchestrator
+from test.utils.unit_test_framework import UnitTestFramework, torch_ref_wrapper
+
+
+def _generate_dup_inputs(bs_slen, dim_size, src_rows, dtype):
+    """Indices sampled WITH replacement -> duplicates within and across tiles."""
+    rng = np.random.RandomState(42)
+    return {
+        "input.must_alias_input": rng.randn(bs_slen, dim_size).astype(dtype),
+        "dim": 0,
+        "index": rng.randint(0, bs_slen, size=src_rows).astype(np.int32),
+        "src": rng.randn(src_rows, dim_size).astype(dtype),
+        "unique_indices": False,
+    }
+
+
+def _output_tensors(kernel_input):
+    return {"output": kernel_input["input.must_alias_input"]}
+
+
+PARAM_NAMES = "bs_slen, dim_size, src_rows, dtype"
+TEST_PARAMS = [
+    (16, 512, 64, nl.float32),      # all dups inside one <128-row tile
+    (16, 512, 64, nl.bfloat16),
+    (10, 256, 200, nl.float32),     # dups crossing the 128-row tile boundary
+    (64, 1024, 512, nl.float32),
+]
+
+
+@pytest_test_metadata(name="ScatterAddDupIndices")
+@pytest_marks(["scatter_add"])
+class TestScatterAddDupIndices:
+    @pytest.mark.fast
+    @pytest_parametrize(PARAM_NAMES, TEST_PARAMS)
+    def test_scatter_add_dup_indices(
+        self,
+        test_manager: Orchestrator,
+        platform_target: Platforms,
+        bs_slen,
+        dim_size,
+        src_rows,
+        dtype,
+    ):
+        is_bf16 = dtype == nl.bfloat16
+        framework = UnitTestFramework(
+            test_manager=test_manager,
+            kernel_entry=scatter_add,
+            torch_ref=torch_ref_wrapper(scatter_add_torch_ref),
+            kernel_input_generator=lambda _: _generate_dup_inputs(bs_slen, dim_size, src_rows, dtype),
+            output_tensor_descriptor=_output_tensors,
+        )
+        framework.run_test(
+            test_config=None,
+            compiler_args=CompilerArgs(
+                platform_target=platform_target,
+                dump_after_lowering=False,
+                logical_nc_config=2,
+            ),
+            atol=1e-1 if is_bf16 else 1e-4,
+            rtol=1e-2 if is_bf16 else 1e-5,
+        )