NCC_IXLV002: Core barrier name mismatch when LNC2 NKI kernel with sendrecv is called 2+ times in a traced graph

[binaryks]

Describe the bug

[Bug] NCC_IXLV002: Core barrier name mismatch when LNC2 NKI kernel with sendrecv is called 2+ times in a traced graph

Environment

• Instance: trn2.3xlarge (LNC2, logical-neuroncore-config: 2)
• OS: Ubuntu 24.04.4 LTS (x86_64)
• neuronx-cc: 2.23.6484.0+3b612583
• torch-neuronx: 2.9.0.2.12.22436+0f1dac25

Description

When an NKI kernel that uses nisa.sendrecv (LNC2 inter-core communication) is launched with grid [2] (i.e., kernel[2](...)) two or more times in a single traced graph, neuronx-cc crashes during compilation with error NCC_IXLV002 (core barrier name mismatch in lnc_verifier).

We first encountered this in a real multi-layer transformer model (Qwen3-MoE with fused QKV kernels using sendrecv for cross-core reduction). The reproducer below is a simplified version that isolates the issue — the out_dim=128 is just one specific configuration that triggers the bug, not the only one. In our production model the graph structure is more complex and also hits this error.

The same graph compiles successfully if the kernel is called only once.

Error Message

[INTERNAL_ERROR] [NCC_IXLV002] Expected core barriers to have name
LocalSendRecv-0_476_sr_cb_begin, but core barrier 6 in subgraph 0
on core 1 has name LocalSendRecv-0_176_sr_cb_begin

Model Name

N/A

Describe the workload type

tensor all reduce

Instance Type

trn2.3xlarge

Release version

• Instance: trn2.3xlarge (LNC2, logical-neuroncore-config: 2)
• OS: Ubuntu 24.04.4 LTS (x86_64)

apt list --installed | grep -i -e neuron

aws-neuronx-collectives/unknown,now 2.30.59.0-f5cdefb39 amd64 [installed]
aws-neuronx-dkms/unknown,now 2.26.5.0 all [installed,upgradable to: 2.26.10.0]
aws-neuronx-oci-hook/unknown,now 2.14.102.0 amd64 [installed]
aws-neuronx-runtime-lib/unknown,now 2.30.51.0-faafe26f0 amd64 [installed]
aws-neuronx-tools/unknown,now 2.28.23.0-f1c114a9d amd64 [installed]

pip list | grep -i -e neuron -e torch -e transformers -e jax

libneuronxla                  2.2.15515.0+50c26cbd
neuronx-cc                    2.23.6484.0+3b612583
neuronx-distributed           0.17.26814+4b18de63
neuronx-distributed-inference 0.8.16251+f3ca5575
torch                         2.9.0
torch-neuronx                 2.9.0.2.12.22436+0f1dac25
torch-xla                     2.9.0
torchvision                   0.24.0
transformers                  4.57.6

Reproduction Steps

Minimal Reproducer

Self-contained script — no external dependencies beyond torch, nki, torch_neuronx:

"""
Minimal reproducer for NCC_IXLV002 compiler bug.

A simple LNC2 NKI kernel that:
  - Each core loads its own row from a [2, H] input (core 0 -> row 0, core 1 -> row 1)
  - sendrecv to exchange data between cores
  - tensor_tensor add to reduce
  - Store result

When this kernel[2] is called 2+ times in a graph followed by a matmul
with output dim <= 128, neuronx-cc crashes with NCC_IXLV002 (core barrier
name mismatch in lnc_verifier).

Environment:
  - neuronx-cc 2.23.6484.0
  - trn2 (LNC2 default)

"""python

import torch
import torch.nn as nn
import nki
import nki.isa as nisa
import nki.language as nl
import torch_neuronx

HIDDEN = 2048


@nki.jit
def lnc2_reduce(x):
    """
    Minimal LNC2 kernel: each core loads its shard, sendrecv + add.
    Input:  x[2, H]   (2 rows, one per core)
    Output: out[1, H]  (sum of both rows)
    """
    H = x.shape[1]

    # Which core am I? (0 or 1)
    core_id = nl.program_id(axis=0)
    other_core = 1 - core_id

    # Each core loads its own row
    my_row = nl.ndarray((1, H), dtype=x.dtype, buffer=nl.sbuf)
    nisa.dma_copy(dst=my_row, src=x[core_id:core_id + 1, :])

    # Receive the other core's row via sendrecv
    recv_buf = nl.ndarray((1, H), dtype=x.dtype, buffer=nl.sbuf)
    nisa.sendrecv(
        src=my_row, dst=recv_buf,
        send_to_rank=other_core, recv_from_rank=other_core,
        pipe_id=0,
    )

    # Reduce: my_row + recv_buf
    nisa.tensor_tensor(my_row, my_row, recv_buf, op=nl.add)

    # Store result
    out = nl.ndarray((1, H), dtype=x.dtype, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=out, src=my_row)
    return out


class ReduceLayer(nn.Module):
    """Wraps the LNC2 kernel + a projection matmul."""

    def __init__(self):
        super().__init__()
        self.proj = nn.Parameter(torch.randn(HIDDEN, HIDDEN, dtype=torch.bfloat16))

    def forward(self, x):
        # x: [1, 1, HIDDEN]
        # Reshape for kernel: [2, HIDDEN] (LNC2 needs 2 rows)
        x_2d = x.view(1, HIDDEN).expand(2, HIDDEN).contiguous()
        # LNC2 kernel call
        reduced = lnc2_reduce[2](x_2d)
        # Project back: [1, HIDDEN] @ [HIDDEN, HIDDEN] -> [1, HIDDEN]
        result = torch.matmul(reduced, self.proj.data)
        return result.unsqueeze(0)  # -> [1, 1, HIDDEN]


class Model(nn.Module):
    def __init__(self, n_layers, out_dim):
        super().__init__()
        self.layers = nn.ModuleList([ReduceLayer() for _ in range(n_layers)])
        self.head = nn.Linear(HIDDEN, out_dim, bias=False).to(torch.bfloat16)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.head(x)


def trace_test(n_layers, out_dim):
    x = torch.randn(1, 1, HIDDEN, dtype=torch.bfloat16)
    model = Model(n_layers, out_dim)
    model.eval()
    try:
        torch_neuronx.trace(
            model, x,
            compiler_args=["--target=trn2", "--model-type", "transformer", "-O1"],
        )
        return "PASS"
    except Exception as e:
        err = str(e)
        if "NCC_IXLV002" in err or "neuronx-cc failed with 70" in err:
            return "NCC_IXLV002"
        return f"other: {type(e).__name__}: {err[:200]}"


def main():
    print(f"{'Config':<45} Result")
    print("-" * 60)

    configs = [
        (1, 128, "1 layer, out=128"),
        (1, 256, "1 layer, out=256"),
        (2, 256, "2 layers, out=256  (PASS expected)"),
        (2, 128, "2 layers, out=128  (FAIL expected)"),
        (2, 2048, "2 layers, out=2048 (PASS expected)"),
        (3, 128, "3 layers, out=128  (FAIL expected)"),
    ]

    for n_layers, out_dim, label in configs:
        r = trace_test(n_layers, out_dim)
        print(f"  {label:<43} {r}")

    print("-" * 60)


if __name__ == "__main__":
    main()

Test Results

Config	Expected	Actual
1 layer, out=128	PASS	✅ PASS
1 layer, out=256	PASS	✅ PASS
2 layers, out=256	PASS	✅ PASS
2 layers, out=128	FAIL	❌ NCC_IXLV002
2 layers, out=2048	PASS	✅ PASS
3 layers, out=128	FAIL	❌ NCC_IXLV002

Trigger Conditions (all must be true)

1. NKI kernel uses nisa.sendrecv (LNC2 inter-core communication)
2. Kernel is launched with grid [2] (LNC2 mode, one program per core)
3. Kernel is called ≥ 2 times in the same traced graph

Regression Issue

• Select this option if this issue appears to be a regression.

Possible Solution

No response

Logs/Context/Additional Information

No response

[github-actions]

Hi @binaryks, Thank you for filing the issue! We will take a look and get back to you.

[KevGomes1403]

I think this bug is fundamentally about having core_barriers to synchronize physical neuron cores with LNC=2. I've reproduced this by changing the sendrecv's to instead synchronize over a shared HBM buffer.

[INTERNAL_ERROR] [NCC_IXLV002] Expected core barriers to have name CoreBarrier-0_480, but core barrier 5 in subgraph 0 on core 1 has name CoreBarrier-0_180

"""
Minimal reproducer for NCC_IXLV002 compiler bug — core_barrier variant.

A simple LNC2 NKI kernel that:
  - Each core loads its own row from a [2, H] input
  - Stores partial to shared HBM
  - core_barrier to sync
  - Loads the other core's partial and adds
  - Writes result

When this kernel[2] is called 2+ times in a graph followed by a matmul
with output dim <= 128, neuronx-cc may crash with NCC_IXLV002 (core barrier
name mismatch in lnc_verifier).

Environment:
  - neuronx-cc 2.23.6484.0
  - trn2 (LNC2 default)
"""

import os
os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2"
os.environ["NEURON_LOGICAL_NC_CONFIG"] = "2"

import torch
import torch.nn as nn
import nki
import nki.isa as nisa
from nki.isa import core_barrier
import nki.language as nl
import torch_neuronx

HIDDEN = 2048
N_PRGS = 2
H_SHARD = HIDDEN // N_PRGS  # = 1024


@nki.jit
def lnc2_allreduce(x):
    """
    Minimal LNC2 kernel: each core loads its row, stores to shared HBM,
    core_barrier, then loads the other core's row and reduces.

    Input:  x[2, H]   — 2 rows, one per core
    Output: (out[2, H_SHARD], workspace[2, H])
      out[core_id, :] = sum of both rows for this core's H_SHARD columns
      workspace is the shared-HBM partial buffer (must be returned)
    """
    H = x.shape[1]
    core_id = nl.program_id(axis=0)
    other = 1 - core_id

    # Each core loads its own row
    my_row = nl.ndarray((1, H), dtype=x.dtype, buffer=nl.sbuf)
    nisa.dma_copy(dst=my_row, src=x[core_id:core_id + 1, :])

    # Shared-HBM workspace: [2, H] — each core writes its own slot
    workspace = nl.ndarray((N_PRGS, H), dtype=x.dtype, buffer=nl.shared_hbm)
    nisa.dma_copy(dst=workspace[core_id:core_id + 1, :], src=my_row)

    # Sync: wait for both cores to finish writing their rows
    core_barrier(workspace, cores=[0, 1])

    # Each core loads its own H_SHARD from both partials, adds, and stores
    out = nl.ndarray((N_PRGS, H_SHARD), dtype=x.dtype, buffer=nl.shared_hbm)

    my_partial   = nl.ndarray((1, H_SHARD), dtype=x.dtype, buffer=nl.sbuf)
    other_partial = nl.ndarray((1, H_SHARD), dtype=x.dtype, buffer=nl.sbuf)
    nisa.dma_copy(
        dst=my_partial,
        src=workspace[core_id:core_id + 1,
                      nl.ds(core_id * H_SHARD, H_SHARD)],
        dge_mode=3,
    )
    nisa.dma_copy(
        dst=other_partial,
        src=workspace[other:other + 1,
                      nl.ds(core_id * H_SHARD, H_SHARD)],
        dge_mode=3,
    )
    nisa.tensor_tensor(my_partial, my_partial, other_partial, op=nl.add)
    nisa.dma_copy(dst=out[core_id:core_id + 1, :], src=my_partial)

    return out, workspace


class AllReduceLayer(nn.Module):
    """Wraps the LNC2 core_barrier kernel + a projection matmul."""

    def __init__(self):
        super().__init__()
        self.proj = nn.Parameter(torch.randn(HIDDEN, HIDDEN, dtype=torch.bfloat16))

    def forward(self, x):
        # x: [1, 1, HIDDEN]
        x_2d = x.view(1, HIDDEN).expand(2, HIDDEN).contiguous()
        out, _ws = lnc2_allreduce[2](x_2d)
        # out: [2, H_SHARD] — reassemble and project
        full = out.view(1, HIDDEN)
        result = torch.matmul(full, self.proj.data)
        return result.unsqueeze(0)  # -> [1, 1, HIDDEN]


class Model(nn.Module):
    def __init__(self, n_layers, out_dim):
        super().__init__()
        self.layers = nn.ModuleList([AllReduceLayer() for _ in range(n_layers)])
        self.head = nn.Linear(HIDDEN, out_dim, bias=False).to(torch.bfloat16)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return self.head(x)


def trace_test(n_layers, out_dim):
    x = torch.randn(1, 1, HIDDEN, dtype=torch.bfloat16)
    model = Model(n_layers, out_dim)
    model.eval()
    try:
        torch_neuronx.trace(
            model, x,
            compiler_args=["--target=trn2", "--model-type", "transformer", "-O1"],
        )
        return "PASS"
    except Exception as e:
        err = str(e)
        if "NCC_IXLV002" in err or "neuronx-cc failed with 70" in err:
            return "NCC_IXLV002"
        return f"other: {type(e).__name__}: {err[:200]}"


def main():
    print(f"{'Config':<45} Result")
    print("-" * 60)

    configs = [
        (1, 128,  "1 layer,  out=128"),
        (1, 256,  "1 layer,  out=256"),
        (2, 256,  "2 layers, out=256"),
        (2, 128,  "2 layers, out=128"),
        (2, 2048, "2 layers, out=2048"),
        (3, 128,  "3 layers, out=128"),
    ]

    for n_layers, out_dim, label in configs:
        r = trace_test(n_layers, out_dim)
        print(f"  {label:<43} {r}")

    print("-" * 60)


if __name__ == "__main__":
    main()

[mrkcath-aws]

Hi @KevGomes1403, thank you for providing this information! We are reviewing this internally and will return with a response this week.

[mrkcath-aws]

Hi @KevGomes1403 we have addressed this in the latest 2.29 release and I was able to run your example successfully. Could you please update and file a new issue if you continue to see problems? Thank you!