Neuron: tiny-model decode region graphs compile and run on a NeuronCore (issues #9, #25)

Fixes found by reducing neuronx-cc NCC_IINL902 crashes to minimal reproducers (crash rates over repeated
compiles, the compiler is nondeterministic on these graphs):
- uint8 byte assembly of narrowed values (KV records, e2m1/e4m3 codes, ue8m0 exponents) crashed or
  silently corrupted bytes: producers keep int32 codes under compile and stores narrow once
  (fused_transform, kv_quant, store_kv_cache; neuron_compile.encode_codes).
- int16 packed scales written through a uint8 view and read back in the same graph: the Neuron kernel
  now receives the int16 storage and packs byte pairs arithmetically (_alloc_quant, return_quant).
- fp8 legalization pass treated plain float->bf16 conversions as fp8 carriers (wrong values + crash).
- Projection stage split further (nullcontext instead of a disabled CP session).

Device (tiny model, regions mode): every decode graph compiles; argmax vs golden short 9/9,
medium 8/9 (step 6 near-tie, margin 0.065), long 9/9, cos 0.997-0.9995.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/model_executor/layers/tilelang_ops/fused_transform.py

@@ -173,7 +173,12 @@ def _alloc_quant(
         else:
             sf = _restore_single_head_axis(sf, sf_shape, input_ndim)
             assert sf.shape == sf_shape and sf.dtype == sf_dtype and sf.is_contiguous()
-        sf_kernel = sf.view(n, num_heads, packs).view(torch.uint8)
+        sf_kernel = sf.view(n, num_heads, packs)
+        if not PLATFORM.is_neuron:
+            sf_kernel = sf_kernel.view(torch.uint8)
+        # NEURON: the int16 storage itself goes to the kernel, whose compiled path packs the ue8m0 byte pairs
+        # arithmetically; a uint8 view of it, read back as int16 in the same graph, crashes neuronx-cc (NCC_IINL902,
+        # issue #25). Its eager path writes the bytes through a uint8 view as before.
     else:
         sf_shape = (n, num_heads) if groups == 1 else (n, num_heads, groups)
         sf_dtype = torch.float8_e4m3fn if cast.sf_format == "e4m3" else torch.float32

deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py

@@ -369,6 +369,15 @@ def fused_norm_rope_hadamard_cast_store_kvcache_tl(num_heads: int, head_dim: int
             osf = out_sf.view(rows, -1)
             if cache_sf_format == "fp32":
                 osf[:, :num_scales].copy_(scale.float())
+            elif torch.compiler.is_compiling() and osf.dtype == torch.int16:
+                # Neuron graphs: join the scale byte pairs into the int16 storage arithmetically (low byte first), as
+                # neuron/norm.py does. Writing them through a uint8 view that the same graph reads back as int16
+                # crashes neuronx-cc (NCC_IINL902, issue #25).
+                e = scale.to(torch.int32)
+                if num_scales % 2:
+                    e = torch.nn.functional.pad(e, (0, 1))
+                v = e[:, 0::2] | (e[:, 1::2] << 8)
+                osf[:, : v.shape[-1]].copy_(torch.where(v >= 32768, v - 65536, v).to(torch.int16))
             else:
                 osf.view(torch.uint8)[:, :num_scales].copy_(scale)
 

deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py

@@ -75,6 +75,8 @@ def make_kv_scale_impl(is_fp4: bool, sf_format: str, use_reciprocal_lut: bool =
         bits = (clamped * _scalar(inv_quant_max, clamped)).view(torch.int32)
         exponent = ((bits - 1) >> 23) + 1
         inverse = ((254 - exponent) << 23).view(torch.float32)
+        if torch.compiler.is_compiling():
+            return exponent, inverse  # int32 until the store (byte assembly in int32, issue #25)
         return exponent.to(torch.uint8), inverse
 
     return compute_scale
@@ -83,12 +85,15 @@ def make_kv_scale_impl(is_fp4: bool, sf_format: str, use_reciprocal_lut: bool =
 def e2m1_codes(x: torch.Tensor) -> torch.Tensor:
     """FP32 -> e2m1 4-bit codes (uint8), round-to-nearest-even, saturating at 6, NaN -> 6."""
     magnitude = torch.nan_to_num(x.abs(), nan=6.0, posinf=6.0).clamp(max=6.0)
+    # int32 codes inside Neuron graphs: narrowing to uint8 before the bytes are assembled crashes neuronx-cc or
+    # corrupts the bytes (NCC_IINL902, issue #25); the store narrows once.
+    dtype = torch.int32 if torch.compiler.is_compiling() else torch.uint8
     code = torch.where(
         magnitude < 2,
         (magnitude * 2).round(),
         torch.where(magnitude < 4, magnitude.round() + 2, (magnitude * 0.5).round() + 4),
-    ).to(torch.uint8)
-    return code | (torch.signbit(x).to(torch.uint8) * 8)
+    ).to(dtype)
+    return code | (torch.signbit(x).to(dtype) * 8)
 
 
 def pack_e2m1(codes: torch.Tensor) -> torch.Tensor:
@@ -99,6 +104,10 @@ def pack_e2m1(codes: torch.Tensor) -> torch.Tensor:
 
 def fp8_e4m3_bytes(x: torch.Tensor) -> torch.Tensor:
     """FP32 -> saturated e4m3 bytes (uint8)."""
+    if torch.compiler.is_compiling():
+        from deepinfer.neuron_ops.neuron_compile import encode_codes
+
+        return encode_codes(torch.clamp(x, -FLOAT8E4M3_MAX, FLOAT8E4M3_MAX))  # int32 codes, see e2m1_codes
     return torch.clamp(x, -FLOAT8E4M3_MAX, FLOAT8E4M3_MAX).to(torch.float8_e4m3fn).view(torch.uint8)
 
 

deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py

@@ -156,7 +156,7 @@ def norm_tl(
             # true division (python-scalar / tensor would lower to reciprocal * scalar)
             inverse = torch.full_like(clamped, FLOAT8E4M3_MAX) / clamped
             sf_values = clamped / FLOAT8E4M3_MAX
-        quantized = fp8_e4m3_bytes(groups * inverse.unsqueeze(-1)).view(torch.float8_e4m3fn)
+        quantized = fp8_e4m3_bytes(groups * inverse.unsqueeze(-1)).to(torch.uint8).view(torch.float8_e4m3fn)
         y.copy_(quantized.view(q_in.shape))
 
         if use_packed_ue8m0 and sf.dtype == torch.int16:

deepinfer/model_executor/layers/tilelang_ops/neuron/store_kv_cache.py

@@ -27,6 +27,7 @@ def _masked_row_store(cache: torch.Tensor, slots: torch.Tensor, values: torch.Te
     safe = torch.where(valid, slots, torch.zeros_like(slots)).long()
     block, offset = safe // block_size, safe % block_size
     old_slot_zero = cache[0, 0].clone()
+    values = values.to(cache.dtype)  # producers keep int32 codes inside Neuron graphs; narrow once here
     cache[block, offset] = values
     # Slot 0 may have received padding rows: use its first valid writer, else restore it.
     hit_zero = valid & (slots == 0)

deepinfer/model_executor/models/v4_modules/attn.py

@@ -8,6 +8,7 @@ draft KV.
 
 from collections.abc import Callable
 
+import contextlib
 import deep_gemm
 import tile_kernels
 import torch
@@ -782,7 +783,8 @@ class V4Attention(V4AttentionBase):
             projected_kv = self.shared_hidden_kv_proj(shared_hidden_bf16)
         neuron_regions.split()
 
-        with userbuffer.cp_session(enabled=use_cp):
+        # Without CP no session is needed; a plain nullcontext lets the Neuron region graphs split inside the block.
+        with userbuffer.cp_session(enabled=use_cp) if use_cp else contextlib.nullcontext():
             # CP: launch the all_gather of the compressor input first, overlapping it with the Q projection below.
             # Under the yoco_full_kv_stored scenario the main branch does not update kv, so this communication step is not needed
             wait_projected_kv = _make_cp_gather_getter(projected_kv, use_cp) if not yoco_full_kv_stored else None
@@ -795,6 +797,7 @@ class V4Attention(V4AttentionBase):
                 shared_hidden_bf16,
             )
             del hidden_states, hidden_states_bf16, shared_hidden_bf16
+            neuron_regions.split()
             kv_swa, wait_kv_swa = self._prepare_swa_kv(
                 kv_swa,
                 kv_cache.swa_cache,
@@ -802,8 +805,10 @@ class V4Attention(V4AttentionBase):
                 use_cp,
             )
 
+            neuron_regions.split()
             q_b, index_q_b = self._project_q_b(q_a, input_metadata.positions)
             del q_a
+            neuron_regions.split()
 
             kv_ca_index = self._compute_kv_ca_index(
                 wait_projected_kv,

deepinfer/neuron_ops/neuron_compile.py

@@ -170,9 +170,11 @@ def legalize_fp8(gm: torch.fx.GraphModule) -> bool:
                 elif dst == FP8 and not is_fp8(src):
                     repl = graph.call_function(_encode, (src,))
                     _swap(node, repl, carriers, carrier=True)
-                else:  # fp8 -> fp8 (copy / device move): bytes
+                elif is_fp8(src) and dst == FP8:  # fp8 -> fp8 (copy / device move): bytes
                     _set_dtype(node, U8)
                     carriers.add(node)
+                else:  # no fp8 involved
+                    continue
                 changed = True
                 continue
             # Creation with dtype=fp8

tests/neuron/compile/test_fp8_legalize.py

@@ -142,3 +142,21 @@ def test_graphs_on_device(case):
             assert (r.int() - g.int()).abs().max() <= 1, case
         else:
             assert torch.equal(r, g), case
+
+
+def test_non_fp8_conversions_untouched():
+    """A bf16/float conversion next to fp8 code must not be treated as an fp8 carrier."""
+
+    def f(x):
+        y = x.to(torch.bfloat16).float()
+        q = (x * 2).to(FP8)
+        return y + 1, q.float()
+
+    x = torch.randn(4, 64)
+    ref = f(x)
+    from deepinfer.neuron_ops.transfer_to_neuron import hide_cuda
+
+    with hide_cuda():
+        got = torch.compile(f, backend=_host_backend(), fullgraph=True, dynamic=False)(x)
+    for r, g in zip(ref, got):
+        assert torch.equal(r, g)