From e846322c0c60cdb0900fc3c12a6e0dd7d7751066 Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 11:08:05 +0800 Subject: [PATCH] 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 --- deepinfer/neuron_ops/nki_gemm.py | 22 ++++++++++++++++++---- deepinfer/neuron_ops/regions.py | 11 ++++++++++- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index 6dc933d..be1af38 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/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: diff --git a/deepinfer/neuron_ops/regions.py b/deepinfer/neuron_ops/regions.py index aa71490..4d05174 100644 --- a/deepinfer/neuron_ops/regions.py +++ b/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)