neuron: fp4 expert GEMMs decode their weight in SBUF on both cores (blockquant.fp4_linear_fused_2core, DEEPINFER_NEURON_FP4_FUSED=1): 512x5120x4608 0.527 -> 0.447 ms, TTFT 825 -> 807 ms; tokens differ by rounding

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

deepinfer/neuron_ops/nki/blockquant.py

@@ -1034,6 +1034,42 @@ def fp4_linear_fused(x, w, scale):
     return out
 
 
+@nki.jit
+def fp4_linear_fused_2core(x, w, scale):
+    """`fp4_linear_fused` with the output columns split across the two physical cores (launch at grid 2), as
+    `fp4_linear_2core` splits `fp4_linear` (DeepInfer, issue #29). Each core stages its own copy of the activation,
+    scale-expands and decodes only its half of the weight in SBUF, and writes its half of the columns."""
+    m, k = x.shape
+    n = w.shape[0]
+    assert w.shape[1] * 2 == k, f"fp4 weight {w.shape} does not pack k={k}"
+    assert k % 64 == 0, f"k={k} must be a multiple of 64 for the uint16 view"
+    programs = nl.num_programs(0)
+    core = nl.program_id(0)
+    assert programs == 2, f"launch grid is {programs}; launch fp4_linear_fused_2core at grid 2"
+    assert n % (programs * PT) == 0, f"n={n} does not divide into {programs} shares of a multiple of {PT}"
+    span = n // programs
+    lo = core * span
+    _check(k, span)
+    mpad = -(-m // PT) * PT
+
+    out = nl.ndarray((mpad, n), dtype=nl.bfloat16, buffer=nl.shared_hbm)
+    x4 = nl.ndarray((mpad, 4, k // 4), dtype=nl.bfloat16, buffer=nl.private_hbm)
+    srep = nl.ndarray((k // 64, 2, 8, span), dtype=nl.uint8, buffer=nl.private_hbm)
+
+    _stage_rows_fp4(x, m, k, x4)
+    _fp4_scale_rows(scale.view(nl.uint16)[lo : lo + span, 0 : k // 64], srep, span, k)
+    _gemm_fp4_fused(
+        x4,
+        w.view(nl.uint16)[lo : lo + span, 0 : k // 4],
+        srep.reshape((k // 4, span)),
+        out[0:mpad, lo : lo + span],
+        mpad,
+        k,
+        span,
+    )
+    return out
+
+
 #: Select byte 1 of each uint16 lane, in place (signed, as `E4M3_SGN2` is).
 FP8_HI_MASK2 = 0xFF00FF00 - (1 << 32)
 #: Select byte 0 of each uint16 lane; a following `<< 8` moves it to byte 1,

deepinfer/neuron_ops/nki_gemm.py

@@ -72,6 +72,11 @@ def prepare_model(model: torch.nn.Module) -> int:
 #: 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"
+#: DEEPINFER_NEURON_FP4_FUSED: fp4 experts decode their weight in SBUF on both cores (blockquant.fp4_linear_fused_2core)
+#: instead of through a bf16 copy in HBM (512 x 5120 x 4608: 0.527 -> 0.447 ms; 128 rows: 0.362 -> 0.245 ms; prompt
+#: chunk TTFT 825 -> 807 ms). Decode's gated experts keep fp4_linear_2core_gated: the fused kernel's input DMA
+#: transposes cannot sit in the gate's dynamic loop (neuronx-cc NCC_IBIR251).
+FP4_FUSED = os.environ.get("DEEPINFER_NEURON_FP4_FUSED", "1") == "1"
 
 
 @functools.cache
@@ -124,6 +129,8 @@ def linear_qdq(x: torch.Tensor, w: torch.Tensor, group: int | None = None) -> to
     and its int16 scales entirely. Returns bf16 ``[m, n]``."""
     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]
+    if w.dtype == torch.int8 and FP4_FUSED and TWO_CORE and weight.shape[0] % 256 == 0 and x.shape[1] % 64 == 0:
+        return _kernel("fp4_linear_fused_2core", 2)(x, weight, scale)[: x.shape[0]]
     name = "fp4_linear" if w.dtype == torch.int8 else "fp8_linear"
     return _linear_kernel(name, weight.shape[0])(x, weight, scale)[: x.shape[0]]