From f7504a995ae9894177bd100073a29c3f534c002f Mon Sep 17 00:00:00 2001 From: qichu Date: Tue, 3 Nov 2026 12:12:32 +0000 Subject: [PATCH] Kimi-K3 MLP: fuse down_proj's E8M0 activation quant into SiTUv2 With ATOM_FP8_BLOCKSCALE_USE_E8M0_SCALE, the SiTUv2 activation fused only a per-token quant, so the shared expert's down_proj (and layer 0's dense one) ran its own per_1x128_e8m0_quant. Add an E8M0 mode to ATOM's SiTUv2 kernel (situ_and_mul_fp8_e8m0_quant) and fuse it. Output is byte-identical to situ_and_mul + per_1x128_e8m0_quant. --- atom/model_ops/kimi_k3/activations.py | 88 +++++++++++++++++-- atom/models/kimi_k3/model.py | 26 ++++-- .../model_ops/test_kimi_k3_situ_e8m0_quant.py | 67 ++++++++++++++ 3 files changed, 165 insertions(+), 16 deletions(-) create mode 100644 tests/model_ops/test_kimi_k3_situ_e8m0_quant.py diff --git a/atom/model_ops/kimi_k3/activations.py b/atom/model_ops/kimi_k3/activations.py index 2125c97f62..1a1cedc572 100644 --- a/atom/model_ops/kimi_k3/activations.py +++ b/atom/model_ops/kimi_k3/activations.py @@ -15,6 +15,8 @@ import triton import triton.language as tl + from atom.model_ops.triton_e8m0_quant import fp8_e8m0_group_scale + _HAS_TRITON = True except ImportError: # pragma: no cover _HAS_TRITON = False @@ -41,6 +43,11 @@ def _situ_and_mul_kernel( inv_linear_beta, HAS_LINEAR: tl.constexpr, BLOCK: tl.constexpr, + s_ptr=None, + stride_s_m=0, + stride_s_g=0, + E8M0_SCALE: tl.constexpr = False, + FP8_MAX: tl.constexpr = 0.0, ): row = tl.program_id(0) col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) @@ -55,6 +62,20 @@ def _situ_and_mul_kernel( if HAS_LINEAR: u = linear_beta * (2.0 * tl.sigmoid(2.0 * u * inv_linear_beta) - 1.0) y = out * u + if E8M0_SCALE: + # quantize the bf16 the unfused path stores; groups past D are skipped + y = y.to(x_ptr.dtype.element_ty).to(tl.float32) + y = y.reshape(BLOCK // 128, 128) + code, multiplier = fp8_e8m0_group_scale(tl.max(tl.abs(y), axis=1), FP8_MAX) + y = y * multiplier[:, None] + y = tl.clamp(y, -FP8_MAX, FP8_MAX, propagate_nan=tl.PropagateNan.ALL) + y = y.reshape(BLOCK) + groups = tl.program_id(1) * (BLOCK // 128) + tl.arange(0, BLOCK // 128) + tl.store( + s_ptr + row * stride_s_m + groups * stride_s_g, + code, + mask=groups < D // 128, + ) tl.store(y_ptr + row * stride_ym + col, y.to(y_ptr.dtype.element_ty), mask=mask) @triton.jit @@ -210,10 +231,16 @@ def situ_and_mul( y = torch.empty((m, d), dtype=x.dtype, device=x.device) if not _HAS_TRITON or m == 0: return _situ_and_mul_torch(x, beta, linear_beta) + _launch_situ_and_mul(x2, y, beta, linear_beta) + return y.reshape(*lead, d) + + +def _launch_situ_and_mul(x2, y, beta, linear_beta, **e8m0_args) -> None: + """_situ_and_mul_kernel over [m, 2*d] ``x2`` into [m, d] ``y``.""" + m, d = y.shape BLOCK = 1024 - grid = (m, triton.cdiv(d, BLOCK)) has_linear = linear_beta is not None - _situ_and_mul_kernel[grid]( + _situ_and_mul_kernel[(m, triton.cdiv(d, BLOCK))]( x2, y, m, @@ -226,8 +253,8 @@ def situ_and_mul( (1.0 / float(linear_beta)) if has_linear else 0.0, HAS_LINEAR=has_linear, BLOCK=BLOCK, + **e8m0_args, ) - return y.reshape(*lead, d) def _situ_and_mul_quant_fake( @@ -275,13 +302,55 @@ def situ_and_mul_quant( return out, scale +def _situ_and_mul_fp8_e8m0_quant_fake( + x: torch.Tensor, beta: float, linear_beta: float | None, transpose_scale: bool +) -> tuple[torch.Tensor, torch.Tensor]: + m, d = x.numel() // x.shape[-1], x.shape[-1] // 2 + out = torch.empty((m, d), dtype=dtypes.fp8, device=x.device) + scale = torch.empty((m, d // 128), dtype=dtypes.fp8_e8m0, device=x.device) + return out, scale + + +# Functional, like per_1x128_e8m0_quant: Inductor cannot functionalize an E8M0 +# mutation. +@torch_compile_guard(gen_fake=_situ_and_mul_fp8_e8m0_quant_fake, mutates_args=[]) +def situ_and_mul_fp8_e8m0_quant( + x: torch.Tensor, beta: float, linear_beta: float | None, transpose_scale: bool +) -> tuple[torch.Tensor, torch.Tensor]: + """``per_1x128_e8m0_quant(situ_and_mul(x, ...))`` in one kernel; d % 128 == 0.""" + x2 = x.reshape(-1, x.shape[-1]) + m, d = x2.shape[0], x2.shape[1] // 2 + assert d % 128 == 0, f"per-1x128 groups need d % 128 == 0, got d={d}" + out = torch.empty((m, d), dtype=dtypes.fp8, device=x.device) + # uint8, viewed as E8M0 below; transpose_scale: column-major, as aiter's + scale = torch.empty( + (d // 128, m) if transpose_scale else (m, d // 128), + dtype=torch.uint8, + device=x.device, + ) + if m: + _launch_situ_and_mul( + x2, + out, + beta, + linear_beta, + s_ptr=scale, + stride_s_m=1 if transpose_scale else d // 128, + stride_s_g=m if transpose_scale else 1, + E8M0_SCALE=True, + FP8_MAX=float(torch.finfo(dtypes.fp8).max), + ) + return out, scale.view(m, d // 128).view(dtypes.fp8_e8m0) + + def situ_and_mul_maybe_quant( x: torch.Tensor, beta: float, linear_beta: float | None, quant: bool = False, + e8m0_scale: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: - """SiTUv2 gated activation, optionally fused with per-token FP8 quant. + """SiTUv2 gated activation, optionally fused with down_proj's FP8 quant. Unifies the quant and non-quant paths behind one call site so the caller always does ``down_proj(x, x_scale=scale)``: @@ -292,10 +361,15 @@ def situ_and_mul_maybe_quant( - ``quant=True``: fuse SiTUv2 + per-token FP8 quant into one aiter kernel and return ``(fp8 [m, d], scale [m, 1])`` for down_proj's ``x_scale=`` path (dense-MLP / shared-expert under ptpc_fp8). - - Unlike :func:`fused_sigmoid_mul_maybe_quant`, there is no scheme selector: the - aiter ``situv2_and_mul_quant`` kernel implements only the per-token scheme. + - ``e8m0_scale``: per-1x128 FP8 with E8M0 scales instead, from ATOM's kernel + (aiter's ``situv2_and_mul_quant`` is per-token only). """ + if quant and e8m0_scale: + from atom.utils import envs + + return situ_and_mul_fp8_e8m0_quant( + x, beta, linear_beta, envs.ATOM_FP8_BLOCKSCALE_WEIGHT_PRESHUFFLE + ) if quant: return situ_and_mul_quant(x, beta, linear_beta) return situ_and_mul(x, beta, linear_beta), None diff --git a/atom/models/kimi_k3/model.py b/atom/models/kimi_k3/model.py index 2a775134b6..11f556dd72 100644 --- a/atom/models/kimi_k3/model.py +++ b/atom/models/kimi_k3/model.py @@ -190,14 +190,15 @@ def __init__( beta: float = 1.0, linear_beta: float | None = None, fused_quant: bool = False, + e8m0_scale: bool = False, ): super().__init__() self.beta = beta self.linear_beta = linear_beta - # Fuse per-token FP8 quant into the activation only when the consuming - # down_proj runs a8w8 per-token FP8 AND linear_beta is set (the aiter - # kernel always applies the linear-beta tanh to the up half). - self.fused_quant = fused_quant and linear_beta is not None + # Fuse down_proj's FP8 quant into the activation. The per-token (aiter) + # kernel always applies the linear-beta tanh, so it needs linear_beta. + self.e8m0_scale = e8m0_scale + self.fused_quant = fused_quant and (e8m0_scale or linear_beta is not None) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: from atom.model_ops.kimi_k3 import situ_and_mul_maybe_quant @@ -205,7 +206,11 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: # Always (activation, scale): (fp8, scale) when fused, else (bf16, None), # so KimiMLP takes the same down_proj(x_scale=) call site. return situ_and_mul_maybe_quant( - x, self.beta, self.linear_beta, quant=self.fused_quant + x, + self.beta, + self.linear_beta, + quant=self.fused_quant, + e8m0_scale=self.e8m0_scale, ) @@ -287,18 +292,21 @@ def __init__( if config.hidden_act != "situ": raise ValueError(f"Unsupported Kimi-K3 activation: {config.hidden_act}") # Fuse SiTUv2 + activation quant when down_proj is a8w8 per-token FP8 - # (ptpc_fp8): the fused kernel emits (fp8, scale) so down_proj skips its - # standalone quant. gate_up/down share a scheme, so probe down_proj. + # (ptpc_fp8) or per-1x128 FP8 with E8M0 scales: the fused kernel emits + # (fp8, scale) so down_proj skips its standalone quant. gate_up/down share + # a scheme, so probe down_proj. down_type, down_dtype = effective_layer_quant( quant_config, f"{prefix}.down_proj" ) - self._fuse_act_quant = ( - down_type == QuantType.per_Token and down_dtype == dtypes.fp8 + down_e8m0_scale = _is_e8m0_blockscale(quant_config, down_type) + self._fuse_act_quant = down_dtype == dtypes.fp8 and ( + down_type == QuantType.per_Token or down_e8m0_scale ) self.act_fn = SituAndMul( beta=getattr(config, "activation_situ_beta", None) or 1.0, linear_beta=getattr(config, "activation_situ_linear_beta", None), fused_quant=self._fuse_act_quant, + e8m0_scale=down_e8m0_scale, ) def forward(self, x: torch.Tensor) -> torch.Tensor: diff --git a/tests/model_ops/test_kimi_k3_situ_e8m0_quant.py b/tests/model_ops/test_kimi_k3_situ_e8m0_quant.py new file mode 100644 index 0000000000..fe15360a57 --- /dev/null +++ b/tests/model_ops/test_kimi_k3_situ_e8m0_quant.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: MIT +"""SiTUv2 with E8M0 scales: byte-identical to situ_and_mul + down_proj's +per_1x128_e8m0_quant.""" + +import pytest +import torch + +pytest.importorskip("aiter") +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("ROCm GPU required", allow_module_level=True) + +from aiter import dtypes + +from atom.model_ops.kimi_k3.activations import ( + situ_and_mul, + situ_and_mul_fp8_e8m0_quant, +) +from atom.model_ops.linear import per_1x128_e8m0_quant + +# Kimi-K3's activation_situ_beta / activation_situ_linear_beta +BETA, LINEAR_BETA = 4.0, 25.0 +# Decode and prefill row counts. +ROWS = [1, 8, 15, 33, 135, 1009, 8065, 16384] + + +def assert_same_bytes(actual, expected): + (q, s), (q_ref, s_ref) = actual, expected + assert q.dtype == q_ref.dtype and s.dtype == s_ref.dtype == dtypes.fp8_e8m0 + assert q.shape == q_ref.shape and s.shape == s_ref.shape + assert s.stride() == s_ref.stride() + assert torch.equal(s.view(torch.uint8), s_ref.view(torch.uint8)) + mismatch = (q.view(torch.uint8) != q_ref.view(torch.uint8)).nonzero() + assert ( + mismatch.numel() == 0 + ), f"{len(mismatch)} fp8 bytes differ, first {mismatch[:4]}" + + +def reference(x, linear_beta, transpose_scale): + d = x.shape[-1] // 2 + act = situ_and_mul(x, BETA, linear_beta).reshape(-1, d) + return per_1x128_e8m0_quant(act, dtypes.fp8, transpose_scale) + + +# d per TP8 rank: shared expert 768, layer-0 dense 4224 (a 1-group last tile). +@pytest.mark.parametrize("transpose_scale", [True, False]) +@pytest.mark.parametrize("d", [768, 4224]) +@pytest.mark.parametrize("rows", ROWS) +def test_situ_e8m0_matches_situ_and_mul_then_down_proj_quant(rows, d, transpose_scale): + torch.manual_seed(20261107) + x = torch.randn(rows, 2 * d, dtype=torch.bfloat16, device="cuda") * 3 + if rows >= 3: + x[0, :128] = 0 # all-zero groups + x[1, 128:256] = 0 + x[2, d + 256 : d + 384] = 0 # up 0 -> an all-zero group + + expected = reference(x, LINEAR_BETA, transpose_scale) + actual = situ_and_mul_fp8_e8m0_quant(x, BETA, LINEAR_BETA, transpose_scale) + assert_same_bytes(actual, expected) + + +def test_situ_e8m0_without_linear_beta(): + torch.manual_seed(20261108) + x = torch.randn(135, 2 * 768, dtype=torch.bfloat16, device="cuda") * 3 + assert_same_bytes( + situ_and_mul_fp8_e8m0_quant(x, BETA, None, True), reference(x, None, True) + )