From 45e1e44444af994e7e021e374119799971b9c0fc Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 22 Oct 2026 18:38:01 +0800 Subject: [PATCH] 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 --- .../layers/tilelang_ops/fused_transform.py | 7 ++++++- .../tilelang_ops/neuron/fused_transform.py | 9 +++++++++ .../layers/tilelang_ops/neuron/kv_quant.py | 13 +++++++++++-- .../layers/tilelang_ops/neuron/norm.py | 2 +- .../tilelang_ops/neuron/store_kv_cache.py | 1 + .../model_executor/models/v4_modules/attn.py | 7 ++++++- deepinfer/neuron_ops/neuron_compile.py | 4 +++- tests/neuron/compile/test_fp8_legalize.py | 18 ++++++++++++++++++ 8 files changed, 55 insertions(+), 6 deletions(-) diff --git a/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py b/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py index 736dc6a..dfb0513 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/fused_transform.py +++ b/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 diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py index 55fdd67..70a5185 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/fused_transform.py +++ b/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) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py index 92ffb34..b968523 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/kv_quant.py +++ b/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) diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py index e5c8740..2bd5771 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/norm.py +++ b/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: diff --git a/deepinfer/model_executor/layers/tilelang_ops/neuron/store_kv_cache.py b/deepinfer/model_executor/layers/tilelang_ops/neuron/store_kv_cache.py index b821874..a37f552 100644 --- a/deepinfer/model_executor/layers/tilelang_ops/neuron/store_kv_cache.py +++ b/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) diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index 1b8a4ee..e20f586 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/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, diff --git a/deepinfer/neuron_ops/neuron_compile.py b/deepinfer/neuron_ops/neuron_compile.py index 5c4c871..1f39433 100644 --- a/deepinfer/neuron_ops/neuron_compile.py +++ b/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 diff --git a/tests/neuron/compile/test_fp8_legalize.py b/tests/neuron/compile/test_fp8_legalize.py index e10a64f..e2bddf4 100644 --- a/tests/neuron/compile/test_fp8_legalize.py +++ b/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)