Neuron: block-quant GEMMs on both physical cores (grid 2); sync experiment knob (#27)

At LNC=2 a logical NeuronCore is two physical cores and a kernel launched at
grid 1 only has code for one of them. The vllm-neuron#71 _2core kernels split
the output columns by program_id; launched at grid 2 they are 1.7-1.9x faster
on every decode GEMM shape and bitwise equal on one core (fp8 1x1280x36864
2.22 -> 1.17 ms, fp4 64x5120x4608 0.68 -> 0.40 ms). Used when n % 256 == 0;
DEEPINFER_NEURON_BQ_2CORE=0 turns it off.

Device profile of a decode attention region (neuron-explorer): 12.9 ms, of
which the q_b_proj fp8 GEMM alone is ~1.8 ms of vector-engine dequant.

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

deepinfer/neuron_ops/nki_gemm.py

@@ -60,13 +60,27 @@ def prepare_model(model: torch.nn.Module) -> int:
     return count
 
 
+#: At LNC=2 a logical NeuronCore is two physical cores, and a kernel launched at grid 1 has code for only one of them:
+#: half the machine idles. The ``_2core`` kernels divide the output columns by ``nl.program_id(0)`` and run at grid 2
+#: (vllm-neuron#71): 1.7-1.9x on every decode GEMM, bitwise equal (issue #27). DEEPINFER_NEURON_BQ_2CORE=0 turns it off.
+TWO_CORE = os.environ.get("DEEPINFER_NEURON_BQ_2CORE", "1") == "1"
+
+
 @functools.cache
-def _kernel(name: str):  # noqa: ANN202
+def _kernel(name: str, grid: int = 0):  # noqa: ANN202
     from torch_neuronx.nki_hop import wrap_nki
 
     from deepinfer.neuron_ops.nki import blockquant
 
-    return wrap_nki(getattr(blockquant, name))
+    kernel = wrap_nki(getattr(blockquant, name))
+    return kernel[grid] if grid else kernel
+
+
+def _linear_kernel(name: str, n: int):  # noqa: ANN202
+    """``name`` at grid 2 when its output columns split into two whole 128-column halves, else at grid 1."""
+    if TWO_CORE and n % 256 == 0:
+        return _kernel(f"{name}_2core", 2)
+    return _kernel(name)
 
 
 def _activation_bf16(a: torch.Tensor, sfa: torch.Tensor, k: int) -> torch.Tensor:
@@ -92,7 +106,7 @@ def gemm(a: torch.Tensor, sfa: torch.Tensor, b: torch.Tensor, d: torch.Tensor, g
     scale = b._nki_scale if group is None else b._nki_scale[group]
     w = b._nki_weight if group is None else b._nki_weight[group]
     name = "fp4_linear" if b.dtype == torch.int8 else "fp8_linear_raw"
-    out = _kernel(name)(x, w, scale)
+    out = _linear_kernel(name, w.shape[0])(x, w, scale)
     d.copy_(out[:m].to(d.dtype))
 
 
@@ -103,7 +117,7 @@ def linear_qdq(x: torch.Tensor, w: torch.Tensor, group: int | None = None) -> to
     scale = w._nki_scale if group is None else w._nki_scale[group]
     weight = w._nki_weight if group is None else w._nki_weight[group]
     name = "fp4_linear" if w.dtype == torch.int8 else "fp8_linear"
-    return _kernel(name)(x, weight, scale)[: x.shape[0]]
+    return _linear_kernel(name, weight.shape[0])(x, weight, scale)[: x.shape[0]]
 
 
 def has_kernel(w: torch.Tensor | None) -> bool:

deepinfer/neuron_ops/regions.py

@@ -8,6 +8,7 @@ fullgraph. Layers of the same kind trace to the same graph, so neuronx-cc's NEFF
 from __future__ import annotations
 
 import functools
+import os
 import time
 from collections import defaultdict
 from collections.abc import Callable
@@ -39,6 +40,9 @@ def raise_dynamo_limits() -> None:
 _ENTRIES: dict[tuple[type, str], Callable[..., Any]] = {}
 #: DEEPINFER_NEURON_REGION_TIMING: seconds per region kind in the current step (each call synchronized).
 _TIMES: dict[str, list[float]] = defaultdict(list)
+#: Experiment (issue #27): synchronize after every N region calls (0 = never).
+_SYNC_EVERY = int(os.environ.get("DEEPINFER_NEURON_SYNC_EVERY", "0"))
+_CALLS = [0]
 
 
 def take_timings() -> dict[str, list[float]]:
@@ -71,7 +75,12 @@ def entry(module: torch.nn.Module, name: str) -> Callable[..., Any]:
         def fn(mod: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
             with hide_cuda():
                 if not envs.DEEPINFER_NEURON_REGION_TIMING:
-                    return compiled(mod, *args, **kwargs)
+                    out = compiled(mod, *args, **kwargs)
+                    if _SYNC_EVERY:
+                        _CALLS[0] += 1
+                        if _CALLS[0] % _SYNC_EVERY == 0:
+                            torch.neuron.synchronize()
+                    return out
                 torch.neuron.synchronize()
                 started = time.perf_counter()
                 out = compiled(mod, *args, **kwargs)