diff --git a/deepinfer/neuron_ops/nki/blockquant.py b/deepinfer/neuron_ops/nki/blockquant.py index e87a786..6787ba5 100644 --- a/deepinfer/neuron_ops/nki/blockquant.py +++ b/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, diff --git a/deepinfer/neuron_ops/nki_gemm.py b/deepinfer/neuron_ops/nki_gemm.py index 62ef83d..f4490e8 100644 --- a/deepinfer/neuron_ops/nki_gemm.py +++ b/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]]