neuron: mHC collapse through torch by default (DEEPINFER_NEURON_HC_TORCH=collapse); fix Neuron bind_cpu core index

collapse in torch is bit-exact with the hc_collapse kernel and a faster graph (512 rows: 0.15 vs 1.73 ms
standalone). post_torch (the kernel's operation order) is bit-equal standalone but rounds differently fused
in the layer graph, so it is opt-in. Tiny golden 9/9 with collapse. bind_cpu receives the physical core id;
_neuron_bind_cpus indexed CUDA_VISIBLE_DEVICES with it again (broke runs on cores other than 0..n).

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

deepinfer/model_executor/models/v4_modules/mhc_ops.py

@@ -39,13 +39,15 @@ def mhc_post(
     post_layer_mix: torch.Tensor,
     comb_res_mix: torch.Tensor,
 ) -> torch.Tensor:
-    if nki_hc.active(x):
+    if nki_hc.active(x, "post"):
         return nki_hc.post(x, residual, post_layer_mix, comb_res_mix)
+    if nki_hc.active(x):  # post routed to torch (DEEPINFER_NEURON_HC_TORCH): the kernel's operation order
+        return nki_hc.post_torch(x, residual, post_layer_mix, comb_res_mix)
     return mhc_post_fwd(x, residual, post_layer_mix, comb_res_mix)
 
 
 def mhc_pre_apply_mix(x: torch.Tensor, mix: torch.Tensor) -> torch.Tensor:
-    if nki_hc.active(x):
+    if nki_hc.active(x, "collapse"):
         return nki_hc.pre_apply_mix(x, mix)
     h = x.shape[-1]
     mhc = mix.shape[-2]

deepinfer/neuron_ops/nki_hc.py

@@ -30,8 +30,15 @@ import torch
 ENABLED = os.environ.get("DEEPINFER_NEURON_NKI_HC", "1") == "1"
 
 
-def active(t: torch.Tensor) -> bool:
-    return ENABLED and t.device.type == "neuron"
+def active(t: torch.Tensor, op: str | None = None) -> bool:
+    return ENABLED and t.device.type == "neuron" and op not in _TORCH_OPS
+
+
+#: DEEPINFER_NEURON_HC_TORCH: mHC ops that take torch even with the NKI kernels enabled. collapse (the default) is
+#: bit-exact and compiles to a faster graph than its kernel. post is opt-in: post_torch is bit-equal to the kernel
+#: standalone, but fused into the layer graph it rounds differently (tiny golden 5/9 on the long prompt). The mixes
+#: keep the kernel path (the shim's fp32 fn GEMM runs on bf16-narrowed operands in a graph, see above).
+_TORCH_OPS = set(filter(None, os.environ.get("DEEPINFER_NEURON_HC_TORCH", "collapse").split(",")))
 
 
 _KERNELS: dict = {}
@@ -80,6 +87,20 @@ def post(x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor,
     )
 
 
+def post_torch(x: torch.Tensor, residual: torch.Tensor, post_layer_mix: torch.Tensor, comb_res_mix: torch.Tensor) -> torch.Tensor:
+    """``post`` in torch, in ``hc_attn_out``'s fp32 operation order (``comb`` terms first, ``post * x`` last), so it
+    compiles to plain elementwise work instead of the kernel (DEEPINFER_NEURON_HC_TORCH=post)."""
+    b, n, hc, h = residual.shape
+    res = residual.float()
+    xf = x.reshape(b, n, 1, h).float()
+    comb = comb_res_mix.reshape(b, n, hc, hc)
+    acc = res[:, :, 0:1, :] * comb[:, :, 0, :].unsqueeze(-1)  # [b, n, j, h]: comb[0, j] * res[0]
+    for i in range(1, hc):
+        acc = res[:, :, i : i + 1, :] * comb[:, :, i, :].unsqueeze(-1) + acc
+    out = xf * post_layer_mix.reshape(b, n, hc, 1) + acc
+    return out.to(residual.dtype)
+
+
 def split_fn(fn: torch.Tensor) -> None:
     """Attach ``fn``'s bf16 ``hi`` / ``lo`` halves (``hi + lo == fn`` to ~2**-17); eager only, see module doc."""
     hi = fn.detach().to(torch.bfloat16)
@@ -150,4 +171,8 @@ def pre_big_fuse(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """``mhc_pre_big_fuse``: ``(post [..., hc, 1], comb [..., hc, hc], layer_input [..., h])``."""
     pre, post_mix, comb = mixes(residual, fn, scale, base)
+    if "collapse" in _TORCH_OPS:
+        from deepinfer.model_executor.models.v4_modules.mhc_ops import mhc_pre_apply_mix
+
+        return post_mix, comb, mhc_pre_apply_mix(residual, pre)
     return post_mix, comb, pre_apply_mix(residual, pre)

deepinfer/utils.py

@@ -141,7 +141,7 @@ def allow_cuda_sync() -> AbstractContextManager[None]:
 
 
 def _neuron_bind_cpus(local_rank: int) -> list[int] | None:
-    """CPUs for the rank on logical NeuronCore ``local_rank`` (of CUDA_VISIBLE_DEVICES), on its device's NUMA node.
+    """CPUs for the rank on logical NeuronCore ``local_rank`` (the physical core id), on its device's NUMA node.
 
     The generic split (physical cores // ranks) gave each of 64 ranks one logical CPU, numbered by rank: the step's
     dispatch thread and the runtime's threads shared it, and devices 0-3 (NUMA 1 on trn2.48xlarge) got CPUs of NUMA 0.
@@ -158,8 +158,7 @@ def _neuron_bind_cpus(local_rank: int) -> list[int] | None:
     except (OSError, subprocess.SubprocessError, ValueError):
         logger.warning("neuron-ls unavailable: generic CPU binding")
         return None
-    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
-    core = int(visible.split(",")[local_rank]) if visible else local_rank
+    core = local_rank  # bind_cpu gets the worker's gpu_id: already the physical core (CUDA_VISIBLE_DEVICES entry)
     device = next((d for d in devices if core in d["neuroncore_ids"]), None)
     if device is None:
         return None