diff --git a/python/sglang/__init__.py b/python/sglang/__init__.py index 3fbec34a6852..f8fa6512fd10 100644 --- a/python/sglang/__init__.py +++ b/python/sglang/__init__.py @@ -13,6 +13,42 @@ _redirect_third_party_caches() +# Kimi-K3 may opt into an SGLang-owned AITER tuning profile. Configure it +# before any downstream import can initialize AITER_CONFIGS. +import importlib.util as _importlib_util +import os as _os +from pathlib import Path as _Path + +import torch as _torch + +if ( + _torch.version.hip is not None + and _os.environ.get("SGLANG_K3_AITER_M16384_PROFILE", "0").lower() in ("1", "true") + and "AITER_CONFIG_GEMM_BF16" not in _os.environ +): + _aiter_spec = _importlib_util.find_spec("aiter") + if _aiter_spec is not None and _aiter_spec.origin is not None: + _aiter_root = _Path(_aiter_spec.origin).resolve().parent + _base = _aiter_root / "configs" / "bf16_tuned_gemm.csv" + _model_configs = sorted( + (_aiter_root / "configs" / "model_configs").glob("*bf16_tuned_gemm*.csv") + ) + _profile = ( + _Path(__file__).resolve().parent + / "kernels" + / "ops" + / "kimi_k3" + / "configs" + / "kimik3_m16384_profile.csv" + ) + _paths = [_base, *_model_configs, _profile] + if all(_path.is_file() for _path in _paths): + _os.environ["AITER_CONFIG_GEMM_BF16"] = _os.pathsep.join(map(str, _paths)) +del _importlib_util +del _os +del _Path +del _torch + if _sys.platform == "darwin" and _platform.machine() == "arm64": from sglang._platform_stubs import install_platform_stubs as _install_platform_stubs diff --git a/python/sglang/kernels/ops/kimi_k3/aiter_fusion.py b/python/sglang/kernels/ops/kimi_k3/aiter_fusion.py new file mode 100644 index 000000000000..a493a25c8f07 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/aiter_fusion.py @@ -0,0 +1,21 @@ +"""Shared helpers for fail-closed Kimi-K3 AITER batch-1 fusions.""" + +from __future__ import annotations + +import torch + +_FP8_MAX = 448.0 + + +def quantize_fp8_rows(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a BF16/FP32 [out, in] weight with one FP32 scale per row.""" + if weight.ndim != 2 or not weight.is_cuda or not weight.is_contiguous(): + raise ValueError("expected a contiguous CUDA [out, in] weight") + row_amax = weight.float().abs().amax(dim=1) + scale = (row_amax / _FP8_MAX).clamp_min(torch.finfo(torch.float32).tiny) + quantized = ( + (weight.float() / scale[:, None]) + .clamp(-_FP8_MAX, _FP8_MAX) + .to(torch.float8_e4m3fn) + ) + return quantized.contiguous(), scale.contiguous() diff --git a/python/sglang/kernels/ops/kimi_k3/configs/kimik3_m16384_profile.csv b/python/sglang/kernels/ops/kimi_k3/configs/kimik3_m16384_profile.csv new file mode 100644 index 000000000000..eb4039c0710b --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/configs/kimik3_m16384_profile.csv @@ -0,0 +1,9 @@ +gfx,cu_num,M,N,K,bias,dtype,outdtype,scaleAB,bpreshuffle,libtype,solidx,splitK,us,kernelName,err_ratio,tflops,bw +gfx950,256,16384,2304,1536,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,438972,0,116.1558,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x224x64_MI16x16x1_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR0_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_7_MO40_NTn1_NTA2_NTB3_NTC7_NTD5_NTM0_NEPBS14_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO1_SVW8_SK3_SKFTR0_SKXCCM8_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB1_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,998.35,1144.21 +gfx950,256,16384,3072,512,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,439111,0,58.0215,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR1_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA2_NTB1_NTC4_NTD2_NTM0_NEPBS12_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO0_SVW8_SK3_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,888.28,2078.3 +gfx950,256,16384,6144,7168,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,439111,0,917.9176,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR1_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA2_NTB1_NTC4_NTD2_NTM0_NEPBS12_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO0_SVW8_SK3_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,1572.16,571.17 +gfx950,256,16384,7168,1536,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,438241,0,285.7495,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR1_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA3_NTB2_NTC4_NTD4_NTM0_NEPBS16_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO0_SVW8_SK3_SKFTR0_SKXCCM4_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,1262.56,1075.18 +gfx950,256,16384,7168,3584,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,439112,0,575.7408,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR1_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA1_NTB3_NTC5_NTD4_NTM0_NEPBS16_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO1_SRVW0_SSO4_SVW8_SK3_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,1462.14,701.19 +gfx950,256,16384,7168,4224,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,439111,0,663.2511,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_CMS_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR1_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSU0_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA2_NTB1_NTC4_NTD2_NTM0_NEPBS12_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO0_SVW8_SK3_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO0_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,1495.87,654.12 +gfx950,256,16384,8448,7168,False,torch.bfloat16,torch.bfloat16,False,False,hipblaslt,439048,0,1400.0434,Cijk_Alik_Bljk_BBS_BH_Bias_HA_S_SAV_UserArgs_MT256x256x64_MI16x16x1_SN_LDSB0_AFC0_AFEM1_AFEM1_ASEM1_CLR0_CADS0_DTLA1_DTLB1_DTVA0_DTVB0_EPS0_FDSI0_GRPM1_GRVWA8_GRVWB8_GSUAMB_GLS0_ISA950_IU1_K1_LDSTI0_LBSPPA1024_LBSPPB1024_LBSPPM0_LPA16_LPB16_LPM0_LRVW8_LWPMn1_MIAV0_MIWT8_8_MO40_NTn1_NTA0_NTB3_NTC4_NTD5_NTM0_NEPBS0_NLCA1_NLCB1_ONLL1_PGR2_PLR1_PKA1_SIA3_SS1_SPO0_SRVW0_SSO0_SVW8_SK0_SKFTR0_SKXCCM0_TLDS1_ULSGRO0_USL1_UIOFGRO0_USFGRO1_VSn1_VWA8_VWB8_WSGRA0_WSGRB0_WS64_WG32_8_1,0.0,1417.3,452.0 +gfx950,256,16384,1536,7168,False,torch.bfloat16,torch.bfloat16,False,False,torch,0,0,297.4444,native,0.0,1212.92,1032.91 diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/__init__.py b/python/sglang/kernels/ops/kimi_k3/flydsl/__init__.py new file mode 100644 index 000000000000..0f8919e5fb65 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/__init__.py @@ -0,0 +1,39 @@ +"""SGLang-maintained Kimi-K3 FlyDSL specializations.""" + +# AITER owns the FlyDSL toolchain bootstrap and shared tensor/buffer shims. +# Import it before local kernel modules so its vendored FlyDSL path is active. +import aiter as _aiter # noqa: F401 + +from .kimi_k3_kda_input_group64 import ( + kimi_k3_kda_input_group64, + quantize_kimi_k3_kda_input_group64, + supports_kimi_k3_kda_input_group64, +) +from .kimi_k3_mla_gate import kimi_k3_mla_gate, supports_kimi_k3_mla_gate +from .kimi_k3_moe_preroute_fp8 import ( + kimi_k3_moe_tri_projection_fp8, + kimi_k3_shared_down_fp8, + supports_kimi_k3_moe_tri_projection_fp8, + supports_kimi_k3_shared_down_fp8, + supports_kimi_k3_shared_down_fp8_weight, +) +from .latent_moe_tail_fp8 import ( + latent_moe_tail_fp8, + quantize_latent_moe_tail_weight, + supports_latent_moe_tail_fp8, +) + +__all__ = [ + "kimi_k3_kda_input_group64", + "kimi_k3_mla_gate", + "kimi_k3_moe_tri_projection_fp8", + "kimi_k3_shared_down_fp8", + "latent_moe_tail_fp8", + "quantize_kimi_k3_kda_input_group64", + "quantize_latent_moe_tail_weight", + "supports_kimi_k3_kda_input_group64", + "supports_kimi_k3_mla_gate", + "supports_kimi_k3_moe_tri_projection_fp8", + "supports_kimi_k3_shared_down_fp8", + "supports_kimi_k3_shared_down_fp8_weight", +] diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/__init__.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_dual_projection_fp8_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_dual_projection_fp8_gfx950.py new file mode 100644 index 000000000000..63e9ef1f70e3 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_dual_projection_fp8_gfx950.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Fixed-shape Kimi-K3 B1 BF16-by-FP8 dual projection for gfx950. + +The activation remains BF16. Each weight row is quantized independently to +OCP FP8 E4M3 and carries one FP32 dequantization scale. This keeps runtime +activation quantization out of the decode graph while halving weight traffic. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from aiter.ops.flydsl.kernels.vector import ReductionOp +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as arith_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import T + +_HIDDEN_SIZE = 7168 +_ROUTED_SIZE = 3584 +_SHARED_UP_SIZE = 1536 +_TOTAL_OUTPUT_SIZE = _ROUTED_SIZE + _SHARED_UP_SIZE +_WAVE_SIZE = 64 +_ELEMENTS_PER_LOAD = 8 + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_b1_dual_projection_fp8_module( + rows_per_wave: int = 3, + cu_count: int = 256, + waves_per_eu: int = 0, + weight_cache_modifier: int = 2, + hidden_to_lds: bool = True, +): + """Build the row-scaled FP8-weight/BF16-activation dual projection.""" + + if rows_per_wave not in (1, 2, 3, 4, 5, 6, 8): + raise ValueError("rows_per_wave must be 1, 2, 3, 4, 5, 6, or 8") + if not 1 <= cu_count <= 256: + raise ValueError("cu_count must be between 1 and 256") + if waves_per_eu < 0: + raise ValueError("waves_per_eu must be non-negative") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + + output_groups = (_TOTAL_OUTPUT_SIZE + rows_per_wave - 1) // rows_per_wave + waves_per_block = min(16, (output_groups + cu_count - 1) // cu_count) + block_threads = waves_per_block * _WAVE_SIZE + groups_per_grid = cu_count * waves_per_block + persistent_iterations = (output_groups + groups_per_grid - 1) // groups_per_grid + hidden_load_iterations = ( + _HIDDEN_SIZE + block_threads * _ELEMENTS_PER_LOAD - 1 + ) // (block_threads * _ELEMENTS_PER_LOAD) + + @fx.struct + class SharedStorage: + hidden: fx.Array[fx.BFloat16, _HIDDEN_SIZE, 16] + + kernel_name = ( + "kimi_k3_b1_dual_projection_bf16_fp8_gfx950" + f"_rpw{rows_per_wave}_cu{cu_count}_wpb{waves_per_block}" + f"_wpe{waves_per_eu}_wcm{weight_cache_modifier}" + f"_hlds{int(hidden_to_lds)}" + ) + + @flyc.kernel( + name=kernel_name, + known_block_size=[block_threads, 1, 1], + ) + def dual_projection_fp8_kernel( + hidden: fx.Pointer, + routed_weight: fx.Pointer, + routed_scale: fx.Pointer, + shared_weight: fx.Pointer, + shared_scale: fx.Pointer, + routed_output: fx.Pointer, + shared_output: fx.Pointer, + ): + i32 = T.i32 + f32 = T.f32 + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE_SIZE, type=i32) + wave = tid // arith.constant(_WAVE_SIZE, type=i32) + + hidden_rsrc = ptr_rsrc(hidden) + routed_weight_rsrc = ptr_rsrc(routed_weight) + routed_scale_rsrc = ptr_rsrc(routed_scale) + shared_weight_rsrc = ptr_rsrc(shared_weight) + shared_scale_rsrc = ptr_rsrc(shared_scale) + routed_output_rsrc = ptr_rsrc(routed_output) + shared_output_rsrc = ptr_rsrc(shared_output) + hidden_lds = fx.SharedAllocator().allocate(SharedStorage).peek().hidden.ptr + + vec2_f32 = T.vec(2, f32) + vec8_bf16 = T.vec(_ELEMENTS_PER_LOAD, T.bf16) + vec8_f32 = T.vec(_ELEMENTS_PER_LOAD, f32) + zero_f32 = arith.constant(0.0, type=f32) + + def load_bf16x8(resource, element_index): + dwords = buffer_ops.buffer_load( + resource, + element_index // arith.constant(2, type=i32), + vec_width=4, + dtype=i32, + ) + return vector.bitcast(vec8_bf16, dwords) + + def load_fp8x8_as_f32(resource, element_index): + packed = buffer_ops.buffer_load( + resource, + element_index // arith.constant(4, type=i32), + vec_width=2, + dtype=i32, + cache_modifier=weight_cache_modifier, + ) + packed = ArithValue(packed) + packed0 = vector.extract( + packed, + static_position=[0], + dynamic_position=[], + ) + packed1 = vector.extract( + packed, + static_position=[1], + dynamic_position=[], + ) + weight0_lo = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed0, + word_sel=False, + ) + weight0_hi = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed0, + word_sel=True, + ) + weight1_lo = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed1, + word_sel=False, + ) + weight1_hi = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed1, + word_sel=True, + ) + weight_lo = weight0_lo.shuffle(weight0_hi, [0, 1, 2, 3]) + weight_hi = weight1_lo.shuffle(weight1_hi, [0, 1, 2, 3]) + return weight_lo.shuffle(weight_hi, [0, 1, 2, 3, 4, 5, 6, 7]) + + def wave_reduce_add(value): + reduced = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(reduced).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE_SIZE, type=i32), + ) + ) + reduced = arith_dialect.AddFOp( + reduced, + peer, + fastmath=fm_fast, + ).result + return reduced + + if const_expr(hidden_to_lds): + for load_iteration in range_constexpr(hidden_load_iterations): + element_index = ( + tid + arith.constant(load_iteration * block_threads, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + can_load = arith.cmpi( + CmpIPredicate.ult, + element_index, + arith.constant(_HIDDEN_SIZE, type=i32), + ) + load_if = scf.IfOp(can_load) + with ir.InsertionPoint(load_if.then_block): + hidden_vector = load_bf16x8(hidden_rsrc, element_index) + fx.ptr_store(hidden_vector, hidden_lds + element_index) + scf.YieldOp([]) + gpu.barrier() + + first_group = ( + ArithValue(gpu.block_idx.x) * arith.constant(waves_per_block, type=i32) + + wave + ) + for persistent_index in range_constexpr(persistent_iterations): + group = first_group + arith.constant( + persistent_index * groups_per_grid, + type=i32, + ) + row_base = group * arith.constant(rows_per_wave, type=i32) + for row_offset in range_constexpr(rows_per_wave): + row = row_base + arith.constant(row_offset, type=i32) + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_TOTAL_OUTPUT_SIZE, type=i32), + ) + row_if = scf.IfOp(row_in_range) + with ir.InsertionPoint(row_if.then_block): + is_routed = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_ROUTED_SIZE, type=i32), + ) + shared_row = row - arith.constant(_ROUTED_SIZE, type=i32) + scale_if = scf.IfOp( + is_routed, + results_=[f32], + has_else=True, + ) + with ir.InsertionPoint(scale_if.then_block): + scale = buffer_ops.buffer_load( + routed_scale_rsrc, + row, + vec_width=1, + dtype=f32, + ) + scf.YieldOp([_raw(scale)]) + with ir.InsertionPoint(scale_if.else_block): + scale = buffer_ops.buffer_load( + shared_scale_rsrc, + shared_row, + vec_width=1, + dtype=f32, + ) + scf.YieldOp([_raw(scale)]) + + local_dot = ArithValue(zero_f32) + for k_iteration in range_constexpr( + _HIDDEN_SIZE // (_WAVE_SIZE * _ELEMENTS_PER_LOAD) + ): + k_element = ( + lane + arith.constant(k_iteration * _WAVE_SIZE, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + if const_expr(hidden_to_lds): + hidden_bf16 = fx.ptr_load( + hidden_lds + k_element, + result_type=vec8_bf16, + ) + else: + hidden_bf16 = load_bf16x8(hidden_rsrc, k_element) + hidden_f32 = ArithValue(hidden_bf16).extf(vec8_f32) + weight_if = scf.IfOp( + is_routed, + results_=[vec8_f32], + has_else=True, + ) + with ir.InsertionPoint(weight_if.then_block): + weight_element = ( + row * arith.constant(_HIDDEN_SIZE, type=i32) + k_element + ) + weight_f32 = load_fp8x8_as_f32( + routed_weight_rsrc, + weight_element, + ) + scf.YieldOp([_raw(weight_f32)]) + with ir.InsertionPoint(weight_if.else_block): + weight_element = ( + shared_row * arith.constant(_HIDDEN_SIZE, type=i32) + + k_element + ) + weight_f32 = load_fp8x8_as_f32( + shared_weight_rsrc, + weight_element, + ) + scf.YieldOp([_raw(weight_f32)]) + local_dot = local_dot + ( + hidden_f32 * ArithValue(weight_if.results[0]) + ).reduce(ReductionOp.ADD, fastmath=fm_fast) + + reduced = ArithValue(wave_reduce_add(local_dot)) * ArithValue( + scale_if.results[0] + ) + is_lane_zero = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(0, type=i32), + ) + write_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(write_if.then_block): + result = arith.trunc_f(T.bf16, _raw(reduced)) + output_if = scf.IfOp( + is_routed, + results_=[], + has_else=True, + ) + with ir.InsertionPoint(output_if.then_block): + buffer_ops.buffer_store(result, routed_output_rsrc, row) + scf.YieldOp([]) + with ir.InsertionPoint(output_if.else_block): + buffer_ops.buffer_store( + result, + shared_output_rsrc, + shared_row, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_dual_projection_fp8( + hidden: fx.Pointer, + routed_weight: fx.Pointer, + routed_scale: fx.Pointer, + shared_weight: fx.Pointer, + shared_scale: fx.Pointer, + routed_output: fx.Pointer, + shared_output: fx.Pointer, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + ctx = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in ctx.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + int(waves_per_eu), + ) + dual_projection_fp8_kernel( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + routed_output, + shared_output, + ).launch( + grid=(cu_count, 1, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch_dual_projection_fp8.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_dual_projection_fp8 diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_kda_input_group64_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_kda_input_group64_gfx950.py new file mode 100644 index 000000000000..5cce0604fae3 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_kda_input_group64_gfx950.py @@ -0,0 +1,372 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Fixed Kimi-K3 TP8 B1 KDA input projection for gfx950. + +The schedule is inherited from the measured row-E4M3 projection. The only +numerical change is a FP32 scale per output row and 64 input columns. FP8 +values are converted to BF16 for ``fdot2``; each sixteen-element partial is +scaled in FP32 and accumulated in FP32. The final store is BF16. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as arith_dialect +from flydsl._mlir.dialects import llvm, scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import T + +_INPUT_FEATURES = 7168 +_PADDED_OUTPUT_FEATURES = 6288 +_LOGICAL_OUTPUT_FEATURES = 6284 +_GROUP_SIZE = 64 +_GROUPS_PER_ROW = _INPUT_FEATURES // _GROUP_SIZE +_WAVE_SIZE = 64 +_ELEMENTS_PER_LOAD = 16 +_K_PER_WAVE_ITERATION = _WAVE_SIZE * _ELEMENTS_PER_LOAD + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_kda_input_group64_module( + *, + num_tokens: int = 1, + rows_per_wave: int = 2, + cu_count: int = 256, + waves_per_eu: int = 0, + weight_cache_modifier: int = 2, + hidden_to_lds: bool = True, +): + """Build the fixed ``[M,7168] @ [6284,7168].T -> [M,6288]`` kernel.""" + + if num_tokens not in (1, 2): + raise ValueError("num_tokens must be 1 or 2") + if rows_per_wave not in (1, 2, 3, 4): + raise ValueError("rows_per_wave must be 1, 2, 3, or 4") + if not 1 <= cu_count <= 256: + raise ValueError("cu_count must be between 1 and 256") + if waves_per_eu < 0: + raise ValueError("waves_per_eu must be non-negative") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + + output_groups = (_PADDED_OUTPUT_FEATURES + rows_per_wave - 1) // rows_per_wave + waves_per_block = min(16, (output_groups + cu_count - 1) // cu_count) + block_threads = waves_per_block * _WAVE_SIZE + groups_per_grid = cu_count * waves_per_block + persistent_iterations = (output_groups + groups_per_grid - 1) // groups_per_grid + hidden_load_iterations = ( + _INPUT_FEATURES + block_threads * _ELEMENTS_PER_LOAD - 1 + ) // (block_threads * _ELEMENTS_PER_LOAD) + + @fx.struct + class SharedStorage: + hidden: fx.Array[fx.BFloat16, _INPUT_FEATURES, 16] + + kernel_name = ( + f"kimi_k3_kda_input_m{num_tokens}_n6288_stored6284_k7168" + f"_e4m3g64_gfx950_rpw{rows_per_wave}_cu{cu_count}" + f"_wpb{waves_per_block}_wpe{waves_per_eu}" + f"_wcm{weight_cache_modifier}_hlds{int(hidden_to_lds)}" + ) + + @flyc.kernel(name=kernel_name, known_block_size=[block_threads, 1, 1]) + def projection_kernel( + hidden: fx.Pointer, + weight: fx.Pointer, + weight_scale: fx.Pointer, + output: fx.Pointer, + ): + i1 = ir.IntegerType.get_signless(1) + i32 = T.i32 + f32 = T.f32 + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE_SIZE, type=i32) + wave = tid // arith.constant(_WAVE_SIZE, type=i32) + hidden_rsrc = ptr_rsrc(hidden) + weight_rsrc = ptr_rsrc(weight) + scale_rsrc = ptr_rsrc(weight_scale) + output_rsrc = ptr_rsrc(output) + token = ArithValue(gpu.block_idx.y) + hidden_token_base = token * arith.constant(_INPUT_FEATURES, type=i32) + output_token_base = token * arith.constant(_PADDED_OUTPUT_FEATURES, type=i32) + hidden_lds = fx.SharedAllocator().allocate(SharedStorage).peek().hidden.ptr + vec2_f32 = T.vec(2, f32) + vec2_bf16 = T.vec(2, T.bf16) + vec16_bf16 = T.vec(_ELEMENTS_PER_LOAD, T.bf16) + zero_f32 = arith.constant(0.0, type=f32) + + def load_bf16x16(resource, element_index): + dwords_lo = ArithValue( + buffer_ops.buffer_load( + resource, + element_index // arith.constant(2, type=i32), + vec_width=4, + dtype=i32, + ) + ) + dwords_hi = ArithValue( + buffer_ops.buffer_load( + resource, + (element_index + arith.constant(_ELEMENTS_PER_LOAD // 2, type=i32)) + // arith.constant(2, type=i32), + vec_width=4, + dtype=i32, + ) + ) + dwords = vector.shuffle(dwords_lo, dwords_hi, list(range(8))) + return vector.bitcast(vec16_bf16, dwords) + + def load_fp8x16_as_f32(resource, element_index): + packed = ArithValue( + buffer_ops.buffer_load( + resource, + element_index // arith.constant(4, type=i32), + vec_width=4, + dtype=i32, + cache_modifier=weight_cache_modifier, + ) + ) + converted = [] + for packed_index in range_constexpr(4): + packed_dword = vector.extract( + packed, + static_position=[packed_index], + dynamic_position=[], + ) + weight_lo = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed_dword, + word_sel=False, + ) + weight_hi = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed_dword, + word_sel=True, + ) + converted.append(vector.shuffle(weight_lo, weight_hi, [0, 1, 2, 3])) + weight_lo = vector.shuffle( + converted[0], + converted[1], + list(range(8)), + ) + weight_hi = vector.shuffle( + converted[2], + converted[3], + list(range(8)), + ) + return vector.shuffle(weight_lo, weight_hi, list(range(16))) + + def wave_reduce_add(value): + reduced = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(reduced).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE_SIZE, type=i32), + ) + ) + reduced = arith_dialect.AddFOp(reduced, peer, fastmath=fm_fast).result + return reduced + + if const_expr(hidden_to_lds): + for load_iteration in range_constexpr(hidden_load_iterations): + element_index = ( + tid + arith.constant(load_iteration * block_threads, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + can_load = arith.cmpi( + CmpIPredicate.ult, + element_index, + arith.constant(_INPUT_FEATURES, type=i32), + ) + load_if = scf.IfOp(can_load) + with ir.InsertionPoint(load_if.then_block): + fx.ptr_store( + load_bf16x16(hidden_rsrc, hidden_token_base + element_index), + hidden_lds + element_index, + ) + scf.YieldOp([]) + gpu.barrier() + + first_group = ( + ArithValue(gpu.block_idx.x) * arith.constant(waves_per_block, type=i32) + + wave + ) + for persistent_index in range_constexpr(persistent_iterations): + group = first_group + arith.constant( + persistent_index * groups_per_grid, type=i32 + ) + row_base = group * arith.constant(rows_per_wave, type=i32) + for row_offset in range_constexpr(rows_per_wave): + row = row_base + arith.constant(row_offset, type=i32) + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_PADDED_OUTPUT_FEATURES, type=i32), + ) + row_if = scf.IfOp(row_in_range) + with ir.InsertionPoint(row_if.then_block): + row_has_weight = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_LOGICAL_OUTPUT_FEATURES, type=i32), + ) + weighted_if = scf.IfOp(row_has_weight, results_=[], has_else=True) + with ir.InsertionPoint(weighted_if.then_block): + local_dot = ArithValue(zero_f32) + for k_iteration in range_constexpr( + _INPUT_FEATURES // _K_PER_WAVE_ITERATION + ): + k_element = ( + lane + + arith.constant(k_iteration * _WAVE_SIZE, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + if const_expr(hidden_to_lds): + hidden_bf16 = fx.ptr_load( + hidden_lds + k_element, + result_type=vec16_bf16, + ) + else: + hidden_bf16 = load_bf16x16( + hidden_rsrc, hidden_token_base + k_element + ) + weight_element = ( + row * arith.constant(_INPUT_FEATURES, type=i32) + + k_element + ) + weight_f32 = load_fp8x16_as_f32(weight_rsrc, weight_element) + weight_bf16 = arith.trunc_f(vec16_bf16, _raw(weight_f32)) + chunk_dot = ArithValue(zero_f32) + for pair_index in range_constexpr(_ELEMENTS_PER_LOAD // 2): + hidden_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + hidden_bf16, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + hidden_bf16, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + weight_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + weight_bf16, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + weight_bf16, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + chunk_dot = ArithValue( + llvm.call_intrinsic( + f32, + "llvm.amdgcn.fdot2.f32.bf16", + [ + hidden_pair, + weight_pair, + _raw(chunk_dot), + arith.constant(False, type=i1), + ], + [], + [], + ) + ) + scale_index = row * arith.constant( + _GROUPS_PER_ROW, type=i32 + ) + k_element // arith.constant(_GROUP_SIZE, type=i32) + scale = ArithValue( + buffer_ops.buffer_load( + scale_rsrc, + scale_index, + vec_width=1, + dtype=f32, + ) + ) + local_dot = local_dot + chunk_dot * scale + + reduced = wave_reduce_add(local_dot) + is_lane_zero = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(0, type=i32), + ) + write_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(write_if.then_block): + result = arith.trunc_f(T.bf16, reduced) + buffer_ops.buffer_store( + result, output_rsrc, output_token_base + row + ) + scf.YieldOp([]) + scf.YieldOp([]) + with ir.InsertionPoint(weighted_if.else_block): + is_lane_zero = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(0, type=i32), + ) + zero_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(zero_if.then_block): + buffer_ops.buffer_store( + arith.trunc_f(T.bf16, zero_f32), + output_rsrc, + output_token_base + row, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch( + hidden: fx.Pointer, + weight: fx.Pointer, + weight_scale: fx.Pointer, + output: fx.Pointer, + stream: fx.Stream, + ): + context = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in context.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, int(waves_per_eu) + ) + projection_kernel(hidden, weight, weight_scale, output).launch( + grid=(cu_count, num_tokens, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": (AITER_FLYDSL_KERNARG_PRELOAD_COUNT), + }, + } + return launch diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_mla_gate_epilogue_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_mla_gate_epilogue_gfx950.py new file mode 100644 index 000000000000..a1cd55e7a7a4 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_mla_gate_epilogue_gfx950.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Fixed-shape gfx950 Kimi-K3 MLA gate projection and epilogue.""" + +import math + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate +from flydsl.expr.typing import T + +_HIDDEN = 7168 +_OUTPUT = 1536 +_WAVE_SIZE = 64 +_ELEMENTS_PER_LOAD = 8 +_K_PER_WAVE_ITERATION = _WAVE_SIZE * _ELEMENTS_PER_LOAD +_K_ITERATIONS = _HIDDEN // _K_PER_WAVE_ITERATION +_LOG2E = math.log2(math.e) + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_mla_gate_module( + rows_per_block: int = 4, + waves_per_eu: int = 2, + weight_cache_modifier: int = 0, +): + """Build the fixed B1 BF16 projection with one wave owning each output row.""" + + if rows_per_block not in (1, 2, 4, 8): + raise ValueError("rows_per_block must be 1, 2, 4, or 8") + if waves_per_eu not in (0, 1, 2, 3, 4): + raise ValueError("waves_per_eu must be between 0 and 4") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + + block_threads = rows_per_block * _WAVE_SIZE + kernel_name = ( + f"kimi_k3_mla_gate_b1_bf16_gfx950_r{rows_per_block}" + f"_wpe{waves_per_eu}_wcm{weight_cache_modifier}" + ) + + @flyc.kernel(name=kernel_name, known_block_size=[block_threads, 1, 1]) + def gate_kernel( + hidden: fx.Pointer, + weight: fx.Pointer, + attention: fx.Pointer, + output: fx.Pointer, + ): + i1 = ir.IntegerType.get_signless(1) + i32 = T.i32 + f32 = T.f32 + vec2_bf16 = T.vec(2, T.bf16) + vec8_bf16 = T.vec(_ELEMENTS_PER_LOAD, T.bf16) + + tid = ArithValue(gpu.thread_idx.x) + lane = tid % fx.Int32(_WAVE_SIZE) + wave = tid // fx.Int32(_WAVE_SIZE) + row = ArithValue(gpu.block_idx.x) * fx.Int32(rows_per_block) + wave + hidden_rsrc = ptr_rsrc(hidden) + weight_rsrc = ptr_rsrc(weight) + attention_rsrc = ptr_rsrc(attention) + output_rsrc = ptr_rsrc(output) + + def load_bf16x8(resource, element_index, cache_modifier=0): + packed = buffer_ops.buffer_load( + resource, + element_index // fx.Int32(2), + vec_width=4, + dtype=i32, + cache_modifier=cache_modifier, + ) + return vector.bitcast(vec8_bf16, packed) + + local_dot = fx.Float32(0.0) + row_weight_base = row * fx.Int32(_HIDDEN) + for k_iteration in range_constexpr(_K_ITERATIONS): + k = lane * fx.Int32(_ELEMENTS_PER_LOAD) + fx.Int32( + k_iteration * _K_PER_WAVE_ITERATION + ) + hidden_values = load_bf16x8(hidden_rsrc, k) + weight_values = load_bf16x8( + weight_rsrc, + row_weight_base + k, + weight_cache_modifier, + ) + for pair_index in range_constexpr(_ELEMENTS_PER_LOAD // 2): + hidden_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + hidden_values, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + hidden_values, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + weight_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + weight_values, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + weight_values, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + local_dot = ArithValue( + llvm.call_intrinsic( + f32, + "llvm.amdgcn.fdot2.f32.bf16", + [ + hidden_pair, + weight_pair, + _raw(local_dot), + arith.constant(False, type=i1), + ], + [], + [], + ) + ) + + dot = local_dot + for offset in (32, 16, 8, 4, 2, 1): + peer = dot.shuffle_xor(fx.Int32(offset), fx.Int32(_WAVE_SIZE)) + dot = dot + peer + + lane_zero = arith.cmpi(CmpIPredicate.eq, lane, fx.Int32(0)) + write_if = scf.IfOp(lane_zero, results_=[], has_else=False) + with ir.InsertionPoint(write_if.then_block): + # Match production rounding: projection->bf16, sigmoid->bf16, + # bf16 attention multiply in fp32-equivalent arithmetic, result->bf16. + projected_bf16 = arith.trunc_f(T.bf16, _raw(dot)) + projected = ArithValue(arith.extf(f32, projected_bf16)) + is_negative = arith.cmpf( + CmpFPredicate.OLT, # codespell:ignore + projected, + fx.Float32(0.0), + ) + magnitude = ArithValue( + arith.select(is_negative, _raw(-projected), _raw(projected)) + ) + exp_neg_abs = fx.math.exp2(-magnitude * fx.Float32(_LOG2E)) + denominator = fx.Float32(1.0) + exp_neg_abs + sigmoid = ArithValue( + arith.select( + is_negative, + _raw(exp_neg_abs / denominator), + _raw(fx.Float32(1.0) / denominator), + ) + ) + sigmoid_bf16 = arith.trunc_f(T.bf16, _raw(sigmoid)) + sigmoid_f32 = ArithValue(arith.extf(f32, sigmoid_bf16)) + attention_bf16 = buffer_ops.buffer_load( + attention_rsrc, + row, + vec_width=1, + dtype=T.bf16, + ) + attention_f32 = ArithValue(arith.extf(f32, attention_bf16)) + result = arith.trunc_f(T.bf16, _raw(sigmoid_f32 * attention_f32)) + buffer_ops.buffer_store(result, output_rsrc, row) + scf.YieldOp([]) + + @flyc.jit + def launch_gate( + hidden: fx.Pointer, + weight: fx.Pointer, + attention: fx.Pointer, + output: fx.Pointer, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + context = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in context.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + waves_per_eu, + ) + gate_kernel(hidden, weight, attention, output).launch( + grid=(_OUTPUT // rows_per_block, 1, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch_gate.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_gate diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_shared_down_fp8_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_shared_down_fp8_gfx950.py new file mode 100644 index 000000000000..b54e22564f04 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_shared_down_fp8_gfx950.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Fused Kimi-K3 B1 SiTU and FP8-weight shared down projection for gfx950.""" + +import math +import os +from pathlib import Path + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from aiter.ops.flydsl.kernels.vector import ReductionOp +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as arith_dialect +from flydsl._mlir.dialects import scf +from flydsl.compiler.extern_link import ExternFunction +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu, range_constexpr +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import T + +_SHARED_INTERMEDIATE_SIZE = 768 +_SHARED_GATE_UP_SIZE = 2 * _SHARED_INTERMEDIATE_SIZE +_HIDDEN_SIZE = 7168 +_WAVE_SIZE = 64 +_ELEMENTS_PER_LOAD = 4 + + +def _find_ocml_bitcode() -> str: + """Find ROCm's device math library without pinning a ROCm release.""" + + roots = [ + Path(value) + for value in ( + os.environ.get("ROCM_PATH"), + os.environ.get("ROCM_HOME"), + "/opt/rocm", + ) + if value + ] + patterns = ( + "amdgcn/bitcode/ocml.bc", + "lib/llvm/lib/clang/*/lib/amdgcn/bitcode/ocml.bc", + ) + for root in roots: + for pattern in patterns: + matches = sorted(root.glob(pattern), reverse=True) + if matches: + return str(matches[0]) + raise RuntimeError("unable to locate ROCm OCML bitcode") + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_b1_shared_down_fp8_module( + num_tokens: int = 1, + rows_per_wave: int = 2, + cu_count: int = 256, + waves_per_eu: int = 0, + weight_cache_modifier: int = 2, + situ_beta: float = 4.0, + situ_linear_beta: float = 25.0, +): + """Build the fused SiTU + row-scaled FP8-weight down projection.""" + + if num_tokens != 1: + raise ValueError("num_tokens must be 1") + if rows_per_wave not in (1, 2, 3, 4, 5, 6, 8): + raise ValueError("rows_per_wave must be 1, 2, 3, 4, 5, 6, or 8") + if not 1 <= cu_count <= 256: + raise ValueError("cu_count must be between 1 and 256") + if waves_per_eu < 0: + raise ValueError("waves_per_eu must be non-negative") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + if ( + not math.isfinite(situ_beta) + or not math.isfinite(situ_linear_beta) + or situ_beta <= 0.0 + or situ_linear_beta <= 0.0 + ): + raise ValueError("SiTU beta values must be finite and positive") + + output_groups = (_HIDDEN_SIZE + rows_per_wave - 1) // rows_per_wave + waves_per_block = min(16, (output_groups + cu_count - 1) // cu_count) + block_threads = waves_per_block * _WAVE_SIZE + groups_per_grid = cu_count * waves_per_block + persistent_iterations = (output_groups + groups_per_grid - 1) // groups_per_grid + ocml_bitcode = _find_ocml_bitcode() + ocml_exp_f32 = ExternFunction( + "__ocml_exp_f32", + ["float32"], + "float32", + is_pure=True, + bitcode_path=ocml_bitcode, + ) + ocml_tanh_f32 = ExternFunction( + "__ocml_tanh_f32", + ["float32"], + "float32", + is_pure=True, + bitcode_path=ocml_bitcode, + ) + + @fx.struct + class SharedStorage: + activated: fx.Array[fx.BFloat16, _SHARED_INTERMEDIATE_SIZE, 16] + + beta_tag = f"{situ_beta:g}".replace(".", "p") + linear_beta_tag = f"{situ_linear_beta:g}".replace(".", "p") + kernel_name = ( + f"kimi_k3_b{num_tokens}_situ_shared_down_bf16_fp8_gfx950" + f"_rpw{rows_per_wave}_cu{cu_count}_wpb{waves_per_block}" + f"_wpe{waves_per_eu}_wcm{weight_cache_modifier}" + f"_sb{beta_tag}_slb{linear_beta_tag}" + ) + + @flyc.kernel( + name=kernel_name, + known_block_size=[block_threads, 1, 1], + ) + def shared_down_fp8_kernel( + gate_up: fx.Pointer, + weight: fx.Pointer, + weight_scale: fx.Pointer, + output: fx.Pointer, + ): + i32 = T.i32 + f32 = T.f32 + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE_SIZE, type=i32) + wave = tid // arith.constant(_WAVE_SIZE, type=i32) + + gate_up_rsrc = ptr_rsrc(gate_up) + weight_rsrc = ptr_rsrc(weight) + scale_rsrc = ptr_rsrc(weight_scale) + output_rsrc = ptr_rsrc(output) + token = ArithValue(gpu.block_idx.y) + gate_up_token_base = token * arith.constant( + 2 * _SHARED_INTERMEDIATE_SIZE, type=i32 + ) + output_token_base = token * arith.constant(_HIDDEN_SIZE, type=i32) + activated_lds = ( + fx.SharedAllocator().allocate(SharedStorage).peek().activated.ptr + ) + + vec2_f32 = T.vec(2, f32) + vec4_bf16 = T.vec(_ELEMENTS_PER_LOAD, T.bf16) + vec4_f32 = T.vec(_ELEMENTS_PER_LOAD, f32) + zero_f32 = arith.constant(0.0, type=f32) + one_f32 = arith.constant(1.0, type=f32) + zero_i32 = arith.constant(0, type=i32) + + def sigmoid(value): + # Match vLLM's production 1 / (1 + expf(-x)) evaluation. + exponent = ocml_exp_f32(-value) + return one_f32 / (one_f32 + exponent) + + def tanh(value): + return ocml_tanh_f32(value) + + def load_fp8x4_as_f32(resource, element_index): + packed = buffer_ops.buffer_load( + resource, + element_index // arith.constant(4, type=i32), + vec_width=1, + dtype=i32, + cache_modifier=weight_cache_modifier, + ) + weight_lo = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed, + word_sel=False, + ) + weight_hi = cvt_pk_f32_fp8( + res=vec2_f32, + src=packed, + word_sel=True, + ) + return weight_lo.shuffle(weight_hi, [0, 1, 2, 3]) + + def wave_reduce_add(value): + reduced = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(reduced).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE_SIZE, type=i32), + ) + ) + reduced = arith_dialect.AddFOp( + reduced, + peer, + fastmath=fm_fast, + ).result + return reduced + + # Every resident block computes the 768-element SiTU vector once and + # retains it in LDS while that block processes its output-row tile. + for activation_iteration in range_constexpr( + (_SHARED_INTERMEDIATE_SIZE + block_threads - 1) // block_threads + ): + element = tid + arith.constant( + activation_iteration * block_threads, type=i32 + ) + valid = arith.cmpi( + CmpIPredicate.ult, + element, + arith.constant(_SHARED_INTERMEDIATE_SIZE, type=i32), + ) + activation_if = scf.IfOp(valid) + with ir.InsertionPoint(activation_if.then_block): + gate_bf16 = buffer_ops.buffer_load( + gate_up_rsrc, + gate_up_token_base + element, + vec_width=1, + dtype=T.bf16, + ) + up_bf16 = buffer_ops.buffer_load( + gate_up_rsrc, + gate_up_token_base + + element + + arith.constant(_SHARED_INTERMEDIATE_SIZE, type=i32), + vec_width=1, + dtype=T.bf16, + ) + gate = ArithValue(arith.extf(f32, gate_bf16)) + up = ArithValue(arith.extf(f32, up_bf16)) + beta = arith.constant(float(situ_beta), type=f32) + inv_beta = arith.constant(1.0 / float(situ_beta), type=f32) + linear_beta = arith.constant(float(situ_linear_beta), type=f32) + inv_linear_beta = arith.constant( + 1.0 / float(situ_linear_beta), + type=f32, + ) + activated = ( + beta + * tanh(gate * inv_beta) + * sigmoid(gate) + * linear_beta + * tanh(up * inv_linear_beta) + ) + fx.ptr_store( + arith.trunc_f(T.bf16, _raw(activated)), + activated_lds + element, + ) + scf.YieldOp([]) + gpu.barrier() + + first_group = ( + ArithValue(gpu.block_idx.x) * arith.constant(waves_per_block, type=i32) + + wave + ) + for persistent_index in range_constexpr(persistent_iterations): + group = first_group + arith.constant( + persistent_index * groups_per_grid, + type=i32, + ) + row_base = group * arith.constant(rows_per_wave, type=i32) + for row_offset in range_constexpr(rows_per_wave): + row = row_base + arith.constant(row_offset, type=i32) + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_HIDDEN_SIZE, type=i32), + ) + row_if = scf.IfOp(row_in_range) + with ir.InsertionPoint(row_if.then_block): + row_scale = buffer_ops.buffer_load( + scale_rsrc, + row, + vec_width=1, + dtype=f32, + ) + local_dot = ArithValue(zero_f32) + for k_iteration in range_constexpr( + _SHARED_INTERMEDIATE_SIZE // (_WAVE_SIZE * _ELEMENTS_PER_LOAD) + ): + k_element = ( + lane + arith.constant(k_iteration * _WAVE_SIZE, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + activated_bf16 = fx.ptr_load( + activated_lds + k_element, + result_type=vec4_bf16, + ) + activated_f32 = ArithValue(activated_bf16).extf(vec4_f32) + weight_element = ( + row * arith.constant(_SHARED_INTERMEDIATE_SIZE, type=i32) + + k_element + ) + weight_f32 = load_fp8x4_as_f32(weight_rsrc, weight_element) + local_dot = local_dot + (activated_f32 * weight_f32).reduce( + ReductionOp.ADD, fastmath=fm_fast + ) + + reduced = ArithValue(wave_reduce_add(local_dot)) * ArithValue( + row_scale + ) + is_lane_zero = arith.cmpi(CmpIPredicate.eq, lane, zero_i32) + write_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(write_if.then_block): + buffer_ops.buffer_store( + arith.trunc_f(T.bf16, _raw(reduced)), + output_rsrc, + output_token_base + row, + ) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_shared_down_fp8( + gate_up: fx.Pointer, + weight: fx.Pointer, + weight_scale: fx.Pointer, + output: fx.Pointer, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + ctx = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in ctx.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + int(waves_per_eu), + ) + shared_down_fp8_kernel( + gate_up, + weight, + weight_scale, + output, + ).launch( + grid=(cu_count, num_tokens, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch_shared_down_fp8.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_shared_down_fp8 diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_tri_projection_fp8_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_tri_projection_fp8_gfx950.py new file mode 100644 index 000000000000..097305823603 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_tri_projection_fp8_gfx950.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Kimi-K3 B1 mixed-precision MoE tri-projection for gfx950. + +The kernel preserves the accepted pre-route projection precision boundaries: + +* routed-down and shared gate/up use row-scaled OCP FP8 E4M3 weights and + materialize BF16 outputs; +* router logits use BF16 weights, the accepted DPP reduction order, a BF16 + rounding boundary, and an FP32 output. + +All three projections consume one BF16 hidden row from LDS in one wide grid. +""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from aiter.ops.flydsl.kernels.vector import ReductionOp +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as arith_dialect +from flydsl._mlir.dialects import llvm, scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import ( + arith, + const_expr, + gpu, + range_constexpr, + rocdl, +) +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import T + +_HIDDEN_SIZE = 7168 +_ROUTED_SIZE = 3584 +_SHARED_UP_SIZE = 1536 +_FP8_OUTPUT_SIZE = _ROUTED_SIZE + _SHARED_UP_SIZE +_ROUTER_SIZE = 896 +_TOTAL_OUTPUT_SIZE = _FP8_OUTPUT_SIZE + _ROUTER_SIZE +_WAVE_SIZE = 64 +_ELEMENTS_PER_LOAD = 8 + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_b1_tri_projection_fp8_module( + num_tokens: int = 1, + rows_per_wave: int = 2, + cu_count: int = 248, + waves_per_eu: int = 0, + weight_cache_modifier: int = 0, + hidden_to_lds: bool = True, + use_fp8_fdot2: bool = True, +): + """Build the fixed mixed-precision tri-projection launcher.""" + + if num_tokens != 1: + raise ValueError("num_tokens must be 1") + if rows_per_wave not in (1, 2, 3, 4, 5, 6, 8): + raise ValueError("rows_per_wave must be 1, 2, 3, 4, 5, 6, or 8") + if not 1 <= cu_count <= 256: + raise ValueError("cu_count must be between 1 and 256") + if waves_per_eu < 0: + raise ValueError("waves_per_eu must be non-negative") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + + output_groups = (_TOTAL_OUTPUT_SIZE + rows_per_wave - 1) // rows_per_wave + waves_per_block = min(16, (output_groups + cu_count - 1) // cu_count) + block_threads = waves_per_block * _WAVE_SIZE + groups_per_grid = cu_count * waves_per_block + persistent_iterations = (output_groups + groups_per_grid - 1) // groups_per_grid + hidden_load_iterations = ( + _HIDDEN_SIZE + block_threads * _ELEMENTS_PER_LOAD - 1 + ) // (block_threads * _ELEMENTS_PER_LOAD) + + @fx.struct + class SharedStorage: + hidden: fx.Array[fx.BFloat16, _HIDDEN_SIZE, 16] + + kernel_name = ( + f"kimi_k3_b{num_tokens}_tri_projection_bf16_fp8_gfx950" + f"_rpw{rows_per_wave}_cu{cu_count}_wpb{waves_per_block}" + f"_wpe{waves_per_eu}_wcm{weight_cache_modifier}" + f"_hlds{int(hidden_to_lds)}" + f"_fdot2{int(use_fp8_fdot2)}" + ) + + @flyc.kernel( + name=kernel_name, + known_block_size=[block_threads, 1, 1], + ) + def tri_projection_fp8_kernel( + hidden: fx.Pointer, + routed_weight: fx.Pointer, + routed_scale: fx.Pointer, + shared_weight: fx.Pointer, + shared_scale: fx.Pointer, + router_weight: fx.Pointer, + routed_output: fx.Pointer, + shared_output: fx.Pointer, + router_output: fx.Pointer, + ): + i32 = T.i32 + f32 = T.f32 + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE_SIZE, type=i32) + wave = tid // arith.constant(_WAVE_SIZE, type=i32) + + hidden_rsrc = ptr_rsrc(hidden) + routed_weight_rsrc = ptr_rsrc(routed_weight) + routed_scale_rsrc = ptr_rsrc(routed_scale) + shared_weight_rsrc = ptr_rsrc(shared_weight) + shared_scale_rsrc = ptr_rsrc(shared_scale) + router_weight_rsrc = ptr_rsrc(router_weight) + routed_output_rsrc = ptr_rsrc(routed_output) + shared_output_rsrc = ptr_rsrc(shared_output) + router_output_rsrc = ptr_rsrc(router_output) + token = ArithValue(gpu.block_idx.y) + hidden_token_base = token * arith.constant(_HIDDEN_SIZE, type=i32) + routed_token_base = token * arith.constant(_ROUTED_SIZE, type=i32) + shared_token_base = token * arith.constant(_SHARED_UP_SIZE, type=i32) + router_token_base = token * arith.constant(_ROUTER_SIZE, type=i32) + hidden_lds = fx.SharedAllocator().allocate(SharedStorage).peek().hidden.ptr + + vec2_f32 = T.vec(2, f32) + vec2_bf16 = T.vec(2, T.bf16) + vec8_bf16 = T.vec(_ELEMENTS_PER_LOAD, T.bf16) + vec8_f32 = T.vec(_ELEMENTS_PER_LOAD, f32) + zero_i32 = arith.constant(0, type=i32) + zero_f32 = arith.constant(0.0, type=f32) + + def load_bf16x8(resource, element_index): + dwords = buffer_ops.buffer_load( + resource, + element_index // arith.constant(2, type=i32), + vec_width=4, + dtype=i32, + ) + return vector.bitcast(vec8_bf16, dwords) + + def load_fp8x8_as_f32(resource, element_index): + packed = buffer_ops.buffer_load( + resource, + element_index // arith.constant(4, type=i32), + vec_width=2, + dtype=i32, + cache_modifier=weight_cache_modifier, + ) + packed = ArithValue(packed) + packed0 = vector.extract(packed, static_position=[0], dynamic_position=[]) + packed1 = vector.extract(packed, static_position=[1], dynamic_position=[]) + weight0_lo = cvt_pk_f32_fp8(res=vec2_f32, src=packed0, word_sel=False) + weight0_hi = cvt_pk_f32_fp8(res=vec2_f32, src=packed0, word_sel=True) + weight1_lo = cvt_pk_f32_fp8(res=vec2_f32, src=packed1, word_sel=False) + weight1_hi = cvt_pk_f32_fp8(res=vec2_f32, src=packed1, word_sel=True) + weight_lo = weight0_lo.shuffle(weight0_hi, [0, 1, 2, 3]) + weight_hi = weight1_lo.shuffle(weight1_hi, [0, 1, 2, 3]) + return weight_lo.shuffle(weight_hi, [0, 1, 2, 3, 4, 5, 6, 7]) + + def wave_reduce_add(value): + reduced = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(reduced).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE_SIZE, type=i32), + ) + ) + reduced = arith_dialect.AddFOp( + reduced, + peer, + fastmath=fm_fast, + ).result + return reduced + + def dot_bf16x8(left, right, accumulator): + dot = _raw(accumulator) + for pair_index in range_constexpr(_ELEMENTS_PER_LOAD // 2): + left_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + left, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + left, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + right_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + right, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + right, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + dot = llvm.call_intrinsic( + f32, + "llvm.amdgcn.fdot2.f32.bf16", + [ + left_pair, + right_pair, + dot, + arith.constant(False, type=ir.IntegerType.get_signless(1)), + ], + [], + [], + ) + return ArithValue(dot) + + if const_expr(hidden_to_lds): + for load_iteration in range_constexpr(hidden_load_iterations): + element_index = ( + tid + arith.constant(load_iteration * block_threads, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + can_load = arith.cmpi( + CmpIPredicate.ult, + element_index, + arith.constant(_HIDDEN_SIZE, type=i32), + ) + load_if = scf.IfOp(can_load) + with ir.InsertionPoint(load_if.then_block): + hidden_vector = load_bf16x8( + hidden_rsrc, hidden_token_base + element_index + ) + fx.ptr_store(hidden_vector, hidden_lds + element_index) + scf.YieldOp([]) + gpu.barrier() + + first_group = ( + ArithValue(gpu.block_idx.x) * arith.constant(waves_per_block, type=i32) + + wave + ) + for persistent_index in range_constexpr(persistent_iterations): + group = first_group + arith.constant( + persistent_index * groups_per_grid, + type=i32, + ) + row_base = group * arith.constant(rows_per_wave, type=i32) + for row_offset in range_constexpr(rows_per_wave): + row = row_base + arith.constant(row_offset, type=i32) + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_TOTAL_OUTPUT_SIZE, type=i32), + ) + row_if = scf.IfOp(row_in_range) + with ir.InsertionPoint(row_if.then_block): + is_fp8 = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_FP8_OUTPUT_SIZE, type=i32), + ) + precision_if = scf.IfOp(is_fp8, has_else=True) + with ir.InsertionPoint(precision_if.then_block): + is_routed = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_ROUTED_SIZE, type=i32), + ) + shared_row = row - arith.constant(_ROUTED_SIZE, type=i32) + scale_if = scf.IfOp( + is_routed, + results_=[f32], + has_else=True, + ) + with ir.InsertionPoint(scale_if.then_block): + scale = buffer_ops.buffer_load( + routed_scale_rsrc, + row, + vec_width=1, + dtype=f32, + ) + scf.YieldOp([_raw(scale)]) + with ir.InsertionPoint(scale_if.else_block): + scale = buffer_ops.buffer_load( + shared_scale_rsrc, + shared_row, + vec_width=1, + dtype=f32, + ) + scf.YieldOp([_raw(scale)]) + + local_dot = ArithValue(zero_f32) + for k_iteration in range_constexpr( + _HIDDEN_SIZE // (_WAVE_SIZE * _ELEMENTS_PER_LOAD) + ): + k_element = ( + lane + + arith.constant(k_iteration * _WAVE_SIZE, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + if const_expr(hidden_to_lds): + hidden_bf16 = fx.ptr_load( + hidden_lds + k_element, + result_type=vec8_bf16, + ) + else: + hidden_bf16 = load_bf16x8( + hidden_rsrc, hidden_token_base + k_element + ) + weight_if = scf.IfOp( + is_routed, + results_=[vec8_f32], + has_else=True, + ) + with ir.InsertionPoint(weight_if.then_block): + weight_element = ( + row * arith.constant(_HIDDEN_SIZE, type=i32) + + k_element + ) + weight_f32 = load_fp8x8_as_f32( + routed_weight_rsrc, + weight_element, + ) + scf.YieldOp([_raw(weight_f32)]) + with ir.InsertionPoint(weight_if.else_block): + weight_element = ( + shared_row * arith.constant(_HIDDEN_SIZE, type=i32) + + k_element + ) + weight_f32 = load_fp8x8_as_f32( + shared_weight_rsrc, + weight_element, + ) + scf.YieldOp([_raw(weight_f32)]) + if const_expr(use_fp8_fdot2): + weight_bf16 = arith.trunc_f( + vec8_bf16, + weight_if.results[0], + ) + local_dot = dot_bf16x8( + hidden_bf16, + weight_bf16, + local_dot, + ) + else: + hidden_f32 = ArithValue(hidden_bf16).extf(vec8_f32) + local_dot = local_dot + ( + hidden_f32 * ArithValue(weight_if.results[0]) + ).reduce(ReductionOp.ADD, fastmath=fm_fast) + + reduced = ArithValue(wave_reduce_add(local_dot)) * ArithValue( + scale_if.results[0] + ) + is_lane_zero = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(0, type=i32), + ) + write_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(write_if.then_block): + result = arith.trunc_f(T.bf16, _raw(reduced)) + output_if = scf.IfOp( + is_routed, + results_=[], + has_else=True, + ) + with ir.InsertionPoint(output_if.then_block): + buffer_ops.buffer_store( + result, + routed_output_rsrc, + routed_token_base + row, + ) + scf.YieldOp([]) + with ir.InsertionPoint(output_if.else_block): + buffer_ops.buffer_store( + result, + shared_output_rsrc, + shared_token_base + shared_row, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + + with ir.InsertionPoint(precision_if.else_block): + router_row = row - arith.constant( + _FP8_OUTPUT_SIZE, + type=i32, + ) + accumulator = ArithValue(zero_f32) + for k_base in range_constexpr(0, _HIDDEN_SIZE, _WAVE_SIZE): + k = lane + arith.constant(k_base, type=i32) + if const_expr(hidden_to_lds): + hidden_bf16 = fx.ptr_load(hidden_lds + k) + else: + hidden_bf16 = buffer_ops.buffer_load( + hidden_rsrc, + hidden_token_base + k, + vec_width=1, + dtype=T.bf16, + ) + hidden_f32 = ArithValue(arith.extf(f32, _raw(hidden_bf16))) + weight_index = ( + router_row * arith.constant(_HIDDEN_SIZE, type=i32) + k + ) + weight_bf16 = buffer_ops.buffer_load( + router_weight_rsrc, + weight_index, + vec_width=1, + dtype=T.bf16, + ) + weight_f32 = ArithValue(arith.extf(f32, weight_bf16)) + accumulator = accumulator + hidden_f32 * weight_f32 + + for dpp_control in (0xB1, 0x4E, 0x141, 0x140, 0x142, 0x143): + remote_i32 = rocdl.update_dpp( + i32, + zero_i32, + arith.bitcast(i32, _raw(accumulator)), + dpp_control, + 0xF, + 0xF, + True, + ) + remote = ArithValue(arith.bitcast(f32, remote_i32)) + accumulator = accumulator + remote + + is_last_lane = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(_WAVE_SIZE - 1, type=i32), + ) + store_if = scf.IfOp(is_last_lane) + with ir.InsertionPoint(store_if.then_block): + rounded = arith.trunc_f(T.bf16, _raw(accumulator)) + projected = arith.extf(f32, rounded) + buffer_ops.buffer_store( + projected, + router_output_rsrc, + router_token_base + router_row, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_tri_projection_fp8( + hidden: fx.Pointer, + routed_weight: fx.Pointer, + routed_scale: fx.Pointer, + shared_weight: fx.Pointer, + shared_scale: fx.Pointer, + router_weight: fx.Pointer, + routed_output: fx.Pointer, + shared_output: fx.Pointer, + router_output: fx.Pointer, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + ctx = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in ctx.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, + int(waves_per_eu), + ) + tri_projection_fp8_kernel( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + routed_output, + shared_output, + router_output, + ).launch( + grid=(cu_count, num_tokens, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch_tri_projection_fp8.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_tri_projection_fp8 diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_tri_projection_multitoken_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_tri_projection_multitoken_gfx950.py new file mode 100644 index 000000000000..d05f52d9a4ba --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/kimi_k3_tri_projection_multitoken_gfx950.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Weight-reusing cooperative preactivated Kimi-K3 tri projection.""" + +import math +import os +from pathlib import Path + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from flydsl._mlir import ir +from flydsl._mlir.dialects import arith as arith_dialect +from flydsl._mlir.dialects import llvm, scf +from flydsl.compiler.extern_link import ExternFunction +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu, range_constexpr, rocdl +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import T + +_HIDDEN = 7168 +_ROUTED = 3584 +_SHARED = 1536 +_FP8_OUT = _ROUTED + _SHARED +_ROUTER = 896 +_TOTAL = _FP8_OUT + _ROUTER +_WAVE = 64 +_VEC = 8 +_LOG2E = 1.4426950408889634 + + +def _find_ocml_bitcode() -> str: + for root in ( + os.environ.get("ROCM_PATH"), + os.environ.get("ROCM_HOME"), + "/opt/rocm", + ): + if not root: + continue + for pattern in ( + "amdgcn/bitcode/ocml.bc", + "lib/llvm/lib/clang/*/lib/amdgcn/bitcode/ocml.bc", + ): + matches = sorted(Path(root).glob(pattern), reverse=True) + if matches: + return str(matches[0]) + raise RuntimeError("unable to locate ROCm OCML bitcode") + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_kimi_k3_multitoken_tri_projection_module( + *, + num_tokens: int, + token_tile: int = 8, + cu_count: int = 248, + waves_per_block: int = 4, + waves_per_eu: int = 0, + weight_cache_modifier: int = 2, + interleaved_shared_pairs: bool = False, + fast_situ: bool = False, + situ_beta: float = 4.0, + situ_linear_beta: float = 25.0, +): + cooperative_preactivate_shared = True + if num_tokens not in (2, 4): + raise ValueError("num_tokens must be 2 or 4") + if token_tile != num_tokens: + raise ValueError("token_tile must match num_tokens") + if waves_per_block not in (4, 8): + raise ValueError("waves_per_block must be 4 or 8") + if ( + not math.isfinite(situ_beta) + or not math.isfinite(situ_linear_beta) + or situ_beta <= 0.0 + or situ_linear_beta <= 0.0 + ): + raise ValueError("SiTU beta values must be finite and positive") + block_threads = waves_per_block * _WAVE + groups_per_grid = cu_count * waves_per_block + persistent_iterations = (_TOTAL + groups_per_grid - 1) // groups_per_grid + hidden_elements = token_tile * _HIDDEN + handoff_elements = (waves_per_block // 2) * token_tile * 2 + hidden_load_iterations = (hidden_elements + block_threads * _VEC - 1) // ( + block_threads * _VEC + ) + + @fx.struct + class SharedStorage: + hidden: fx.Array[fx.BFloat16, hidden_elements, 16] + handoff: fx.Array[fx.BFloat16, handoff_elements, 16] + + kernel_name = ( + f"kimi_k3_m{num_tokens}_mixed_tri_bf16_fp8_gfx950" + f"_tt{token_tile}_cu{cu_count}_wpb{waves_per_block}" + f"_wpe{waves_per_eu}_wcm{weight_cache_modifier}" + f"_cooppreact768_fast{int(fast_situ)}" + f"_interleaved{int(interleaved_shared_pairs)}" + ) + ocml_bitcode = _find_ocml_bitcode() + ocml_exp_f32 = ExternFunction( + "__ocml_exp_f32", + ["float32"], + "float32", + is_pure=True, + bitcode_path=ocml_bitcode, + ) + ocml_tanh_f32 = ExternFunction( + "__ocml_tanh_f32", + ["float32"], + "float32", + is_pure=True, + bitcode_path=ocml_bitcode, + ) + + @flyc.kernel(name=kernel_name, known_block_size=[block_threads, 1, 1]) + def kernel( + hidden: fx.Pointer, + routed_weight: fx.Pointer, + routed_scale: fx.Pointer, + shared_weight: fx.Pointer, + shared_scale: fx.Pointer, + router_weight: fx.Pointer, + routed_output: fx.Pointer, + shared_output: fx.Pointer, + router_output: fx.Pointer, + ): + i32 = T.i32 + f32 = T.f32 + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE, type=i32) + wave = tid // arith.constant(_WAVE, type=i32) + hidden_rsrc = ptr_rsrc(hidden) + routed_weight_rsrc = ptr_rsrc(routed_weight) + routed_scale_rsrc = ptr_rsrc(routed_scale) + shared_weight_rsrc = ptr_rsrc(shared_weight) + shared_scale_rsrc = ptr_rsrc(shared_scale) + router_weight_rsrc = ptr_rsrc(router_weight) + routed_output_rsrc = ptr_rsrc(routed_output) + shared_output_rsrc = ptr_rsrc(shared_output) + router_output_rsrc = ptr_rsrc(router_output) + shared_storage = fx.SharedAllocator().allocate(SharedStorage).peek() + hidden_lds = shared_storage.hidden.ptr + handoff_lds = shared_storage.handoff.ptr + vec2_f32 = T.vec(2, f32) + vec2_bf16 = T.vec(2, T.bf16) + vec8_bf16 = T.vec(_VEC, T.bf16) + vec8_f32 = T.vec(_VEC, f32) + zero_i32 = arith.constant(0, type=i32) + zero_f32 = arith.constant(0.0, type=f32) + one_f32 = arith.constant(1.0, type=f32) + beta_f32 = arith.constant(float(situ_beta), type=f32) + inv_beta_f32 = arith.constant(1.0 / float(situ_beta), type=f32) + linear_beta_f32 = arith.constant(float(situ_linear_beta), type=f32) + inv_linear_beta_f32 = arith.constant(1.0 / float(situ_linear_beta), type=f32) + + def sigmoid(value): + if const_expr(fast_situ): + exponent = fx.math.exp2(-value * fx.Float32(_LOG2E)) + else: + exponent = ocml_exp_f32(-value) + return one_f32 / (one_f32 + exponent) + + def tanh(value): + if const_expr(fast_situ): + return 2.0 * sigmoid(2.0 * value) - 1.0 + return ocml_tanh_f32(value) + + def situ(gate, up): + return ( + beta_f32 + * tanh(gate * inv_beta_f32) + * sigmoid(gate) + * linear_beta_f32 + * tanh(up * inv_linear_beta_f32) + ) + + def load_bf16x8(resource, element_index): + packed = buffer_ops.buffer_load( + resource, + element_index // arith.constant(2, type=i32), + vec_width=4, + dtype=i32, + ) + return vector.bitcast(vec8_bf16, packed) + + def load_fp8x8(resource, element_index): + packed = ArithValue( + buffer_ops.buffer_load( + resource, + element_index // arith.constant(4, type=i32), + vec_width=2, + dtype=i32, + cache_modifier=weight_cache_modifier, + ) + ) + p0 = vector.extract(packed, static_position=[0], dynamic_position=[]) + p1 = vector.extract(packed, static_position=[1], dynamic_position=[]) + lo0 = cvt_pk_f32_fp8(res=vec2_f32, src=p0, word_sel=False) + hi0 = cvt_pk_f32_fp8(res=vec2_f32, src=p0, word_sel=True) + lo1 = cvt_pk_f32_fp8(res=vec2_f32, src=p1, word_sel=False) + hi1 = cvt_pk_f32_fp8(res=vec2_f32, src=p1, word_sel=True) + return lo0.shuffle(hi0, [0, 1, 2, 3]).shuffle( + lo1.shuffle(hi1, [0, 1, 2, 3]), list(range(8)) + ) + + def dot2x8(left, right, accumulator): + dot = _raw(accumulator) + for pair in range_constexpr(4): + lp = vector.from_elements( + vec2_bf16, + [ + vector.extract( + left, static_position=[pair * 2], dynamic_position=[] + ), + vector.extract( + left, static_position=[pair * 2 + 1], dynamic_position=[] + ), + ], + ) + rp = vector.from_elements( + vec2_bf16, + [ + vector.extract( + right, static_position=[pair * 2], dynamic_position=[] + ), + vector.extract( + right, static_position=[pair * 2 + 1], dynamic_position=[] + ), + ], + ) + dot = llvm.call_intrinsic( + f32, + "llvm.amdgcn.fdot2.f32.bf16", + [ + lp, + rp, + dot, + arith.constant(False, type=ir.IntegerType.get_signless(1)), + ], + [], + [], + ) + return ArithValue(dot) + + def wave_reduce(value): + result = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(result).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE, type=i32), + ) + ) + result = arith_dialect.AddFOp(result, peer, fastmath=fm_fast).result + return ArithValue(result) + + first_group = ( + ArithValue(gpu.block_idx.x) * arith.constant(waves_per_block, type=i32) + + wave + ) + for tile_start in range_constexpr(0, num_tokens, token_tile): + for load_iter in range_constexpr(hidden_load_iterations): + element = ( + tid + arith.constant(load_iter * block_threads, type=i32) + ) * arith.constant(_VEC, type=i32) + can_load = arith.cmpi( + CmpIPredicate.ult, + element, + arith.constant(hidden_elements, type=i32), + ) + load_if = scf.IfOp(can_load) + with ir.InsertionPoint(load_if.then_block): + token_local = element // arith.constant(_HIDDEN, type=i32) + hidden_col = element % arith.constant(_HIDDEN, type=i32) + source = ( + arith.constant(tile_start, type=i32) + token_local + ) * arith.constant(_HIDDEN, type=i32) + hidden_col + fx.ptr_store( + load_bf16x8(hidden_rsrc, source), + hidden_lds + element, + ) + scf.YieldOp([]) + gpu.barrier() + + for persistent_index in range_constexpr(persistent_iterations): + group = first_group + arith.constant( + persistent_index * groups_per_grid, type=i32 + ) + row = group + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_TOTAL, type=i32), + ) + row_if = scf.IfOp(row_in_range) + with ir.InsertionPoint(row_if.then_block): + is_fp8 = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_FP8_OUT, type=i32), + ) + precision_if = scf.IfOp(is_fp8, has_else=True) + with ir.InsertionPoint(precision_if.then_block): + is_routed = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_ROUTED, type=i32), + ) + shared_row = row - arith.constant(_ROUTED, type=i32) + shared_pair = shared_row // arith.constant(2, type=i32) + shared_role = shared_row % arith.constant(2, type=i32) + if const_expr(interleaved_shared_pairs): + shared_weight_row = shared_row + else: + shared_weight_row = ( + shared_pair + + shared_role * arith.constant(_SHARED // 2, type=i32) + ) + scale_if = scf.IfOp(is_routed, results_=[f32], has_else=True) + with ir.InsertionPoint(scale_if.then_block): + scale = buffer_ops.buffer_load( + routed_scale_rsrc, row, vec_width=1, dtype=f32 + ) + scf.YieldOp([_raw(scale)]) + with ir.InsertionPoint(scale_if.else_block): + scale = buffer_ops.buffer_load( + shared_scale_rsrc, + shared_weight_row, + vec_width=1, + dtype=f32, + ) + scf.YieldOp([_raw(scale)]) + accumulators = [ArithValue(zero_f32) for _ in range(token_tile)] + for k_iter in range_constexpr(_HIDDEN // (_WAVE * _VEC)): + k_element = ( + lane + arith.constant(k_iter * _WAVE, type=i32) + ) * arith.constant(_VEC, type=i32) + weight_if = scf.IfOp( + is_routed, results_=[vec8_f32], has_else=True + ) + with ir.InsertionPoint(weight_if.then_block): + weight = load_fp8x8( + routed_weight_rsrc, + row * arith.constant(_HIDDEN, type=i32) + k_element, + ) + scf.YieldOp([_raw(weight)]) + with ir.InsertionPoint(weight_if.else_block): + weight = load_fp8x8( + shared_weight_rsrc, + shared_weight_row + * arith.constant(_HIDDEN, type=i32) + + k_element, + ) + scf.YieldOp([_raw(weight)]) + weight_bf16 = arith.trunc_f(vec8_bf16, weight_if.results[0]) + for token_local in range_constexpr(token_tile): + h = fx.ptr_load( + hidden_lds + + arith.constant(token_local * _HIDDEN, type=i32) + + k_element, + result_type=vec8_bf16, + ) + accumulators[token_local] = dot2x8( + h, weight_bf16, accumulators[token_local] + ) + lane_zero = arith.cmpi( + CmpIPredicate.eq, lane, arith.constant(0, type=i32) + ) + for token_local in range_constexpr(token_tile): + reduced = wave_reduce(accumulators[token_local]) + reduced = reduced * ArithValue(scale_if.results[0]) + store_if = scf.IfOp(lane_zero) + with ir.InsertionPoint(store_if.then_block): + token = tile_start + token_local + result = arith.trunc_f(T.bf16, _raw(reduced)) + output_if = scf.IfOp( + is_routed, results_=[], has_else=True + ) + with ir.InsertionPoint(output_if.then_block): + buffer_ops.buffer_store( + result, + routed_output_rsrc, + arith.constant(token * _ROUTED, type=i32) + row, + ) + scf.YieldOp([]) + with ir.InsertionPoint(output_if.else_block): + pair_local = wave // arith.constant(2, type=i32) + scratch_index = ( + pair_local + * arith.constant(token_tile * 2, type=i32) + + arith.constant(token_local * 2, type=i32) + + shared_role + ) + fx.ptr_store( + result, + handoff_lds + scratch_index, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + + with ir.InsertionPoint(precision_if.else_block): + router_row = row - arith.constant(_FP8_OUT, type=i32) + accumulators = [ArithValue(zero_f32) for _ in range(token_tile)] + for k_iter in range_constexpr(_HIDDEN // (_WAVE * _VEC)): + k = ( + lane + arith.constant(k_iter * _WAVE, type=i32) + ) * arith.constant(_VEC, type=i32) + weight_bf16 = load_bf16x8( + router_weight_rsrc, + router_row * arith.constant(_HIDDEN, type=i32) + k, + ) + for token_local in range_constexpr(token_tile): + hidden_bf16 = fx.ptr_load( + hidden_lds + + arith.constant(token_local * _HIDDEN, type=i32) + + k, + result_type=vec8_bf16, + ) + accumulators[token_local] = dot2x8( + hidden_bf16, + weight_bf16, + accumulators[token_local], + ) + last_lane = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(_WAVE - 1, type=i32), + ) + for token_local in range_constexpr(token_tile): + accumulator = accumulators[token_local] + for dpp in (0xB1, 0x4E, 0x141, 0x140, 0x142, 0x143): + remote_i32 = rocdl.update_dpp( + i32, + zero_i32, + arith.bitcast(i32, _raw(accumulator)), + dpp, + 0xF, + 0xF, + True, + ) + accumulator = accumulator + ArithValue( + arith.bitcast(f32, remote_i32) + ) + store_if = scf.IfOp(last_lane) + with ir.InsertionPoint(store_if.then_block): + rounded = arith.trunc_f(T.bf16, _raw(accumulator)) + buffer_ops.buffer_store( + arith.extf(f32, rounded), + router_output_rsrc, + arith.constant( + (tile_start + token_local) * _ROUTER, + type=i32, + ) + + router_row, + ) + scf.YieldOp([]) + scf.YieldOp([]) + scf.YieldOp([]) + if const_expr(cooperative_preactivate_shared): + block_first_group = ArithValue(gpu.block_idx.x) * arith.constant( + waves_per_block, type=i32 + ) + arith.constant(persistent_index * groups_per_grid, type=i32) + shared_lower = arith.cmpi( + CmpIPredicate.uge, + block_first_group, + arith.constant(_ROUTED, type=i32), + ) + shared_upper = arith.cmpi( + CmpIPredicate.ult, + block_first_group, + arith.constant(_FP8_OUT, type=i32), + ) + shared_iteration = arith.andi(shared_lower, shared_upper) + shared_iteration_if = scf.IfOp(shared_iteration) + with ir.InsertionPoint(shared_iteration_if.then_block): + gpu.barrier() + is_gate_wave = arith.cmpi( + CmpIPredicate.eq, + wave % arith.constant(2, type=i32), + arith.constant(0, type=i32), + ) + gate_wave_if = scf.IfOp(is_gate_wave) + with ir.InsertionPoint(gate_wave_if.then_block): + pair_local = wave // arith.constant(2, type=i32) + shared_pair = ( + row - arith.constant(_ROUTED, type=i32) + ) // arith.constant(2, type=i32) + for token_local in range_constexpr(token_tile): + scratch_base = pair_local * arith.constant( + token_tile * 2, type=i32 + ) + arith.constant(token_local * 2, type=i32) + gate = ArithValue( + arith.extf( + f32, + _raw( + fx.ptr_load( + handoff_lds + scratch_base, + result_type=T.bf16, + ) + ), + ) + ) + up = ArithValue( + arith.extf( + f32, + _raw( + fx.ptr_load( + handoff_lds + + scratch_base + + arith.constant(1, type=i32), + result_type=T.bf16, + ) + ), + ) + ) + activated = arith.trunc_f(T.bf16, _raw(situ(gate, up))) + buffer_ops.buffer_store( + activated, + shared_output_rsrc, + arith.constant( + (tile_start + token_local) * (_SHARED // 2), + type=i32, + ) + + shared_pair, + ) + scf.YieldOp([]) + gpu.barrier() + scf.YieldOp([]) + gpu.barrier() + + @flyc.jit + def launch( + hidden: fx.Pointer, + routed_weight: fx.Pointer, + routed_scale: fx.Pointer, + shared_weight: fx.Pointer, + shared_scale: fx.Pointer, + router_weight: fx.Pointer, + routed_output: fx.Pointer, + shared_output: fx.Pointer, + router_output: fx.Pointer, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + ctx = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in ctx.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, int(waves_per_eu) + ) + kernel( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + routed_output, + shared_output, + router_output, + ).launch( + grid=(cu_count, 1, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + } + } + return launch + + +__all__ = ["build_kimi_k3_multitoken_tri_projection_module"] diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/latent_moe_tail_fp8_gfx950.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/latent_moe_tail_fp8_gfx950.py new file mode 100644 index 000000000000..ff93e7ce8944 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kernels/latent_moe_tail_fp8_gfx950.py @@ -0,0 +1,369 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Persistent-wave Kimi-K3 B1 RMSNorm, FP8 GEMV, and shared-add kernel.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from aiter.ops.flydsl.kernels import buffer_ops, vector +from aiter.ops.flydsl.kernels.tensor_shim import ( + AITER_FLYDSL_KERNARG_PRELOAD, + AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + ptr_rsrc, +) +from aiter.ops.flydsl.kernels.vector import ReductionOp +from flydsl._mlir import ir +from flydsl._mlir.dialects import llvm, scf +from flydsl.compiler.kernel_function import CompilationContext +from flydsl.expr import arith, const_expr, gpu +from flydsl.expr import math as fmath +from flydsl.expr import range_constexpr +from flydsl.expr.arith import ArithValue, CmpIPredicate +from flydsl.expr.rocdl import cvt_pk_f32_fp8 +from flydsl.expr.typing import T + +_LATENT_DIM = 3584 +_HIDDEN_DIM = 7168 +_WAVE_SIZE = 64 +_ELEMENTS_PER_LOAD = 8 +_NORMALIZE_THREADS = _LATENT_DIM // _ELEMENTS_PER_LOAD +_NORMALIZE_WAVES = _NORMALIZE_THREADS // _WAVE_SIZE +_K_PER_WAVE_ITERATION = _WAVE_SIZE * _ELEMENTS_PER_LOAD + + +def _raw(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def build_b1_latent_moe_tail_fp8_persistent_module( + rows_per_wave: int = 2, + cu_count: int = 256, + waves_per_eu: int = 0, + weight_cache_modifier: int = 2, +): + """Build a block-normalize/one-wave-per-output-group specialization.""" + + if rows_per_wave not in (1, 2, 3, 4): + raise ValueError("rows_per_wave must be 1, 2, 3, or 4") + if not 1 <= cu_count <= 256: + raise ValueError("cu_count must be between 1 and 256") + if waves_per_eu < 0: + raise ValueError("waves_per_eu must be non-negative") + if weight_cache_modifier not in (0, 1, 2, 3): + raise ValueError("weight_cache_modifier must be between 0 and 3") + + output_groups = (_HIDDEN_DIM + rows_per_wave - 1) // rows_per_wave + waves_per_block = min(16, (output_groups + cu_count - 1) // cu_count) + block_threads = waves_per_block * _WAVE_SIZE + if block_threads < _NORMALIZE_THREADS: + raise ValueError("schedule needs at least 448 threads for fixed RMS order") + groups_per_grid = cu_count * waves_per_block + persistent_iterations = (output_groups + groups_per_grid - 1) // groups_per_grid + + @fx.struct + class SharedStorage: + hidden: fx.Array[fx.BFloat16, _LATENT_DIM, 16] + rms_sums: fx.Array[fx.Float32, _NORMALIZE_WAVES, 16] + inverse_rms: fx.Array[fx.Float32, 1, 16] + + kernel_name = ( + f"latent_moe_tail_b1_bf16_fp8_persistent_gfx950" + f"_rpw{rows_per_wave}_cu{cu_count}_wpb{waves_per_block}" + f"_wpe{waves_per_eu}_wcm{weight_cache_modifier}" + ) + + @flyc.kernel(name=kernel_name, known_block_size=[block_threads, 1, 1]) + def tail_kernel( + routed: fx.Pointer, + shared: fx.Pointer, + rms_weight: fx.Pointer, + up_weight: fx.Pointer, + up_scale: fx.Pointer, + output: fx.Pointer, + epsilon: fx.Float32, + ): + i1 = ir.IntegerType.get_signless(1) + i32 = T.i32 + f32 = T.f32 + fm_fast = arith.FastMathFlags.fast + tid = ArithValue(gpu.thread_idx.x) + lane = tid % arith.constant(_WAVE_SIZE, type=i32) + wave = tid // arith.constant(_WAVE_SIZE, type=i32) + + routed_rsrc = ptr_rsrc(routed) + shared_rsrc = ptr_rsrc(shared) + rms_weight_rsrc = ptr_rsrc(rms_weight) + up_weight_rsrc = ptr_rsrc(up_weight) + up_scale_rsrc = ptr_rsrc(up_scale) + output_rsrc = ptr_rsrc(output) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + hidden_lds = lds.hidden.ptr + rms_sums = lds.rms_sums.ptr + inverse_rms = lds.inverse_rms.ptr + + vec2_f32 = T.vec(2, f32) + vec2_bf16 = T.vec(2, T.bf16) + vec8_bf16 = T.vec(_ELEMENTS_PER_LOAD, T.bf16) + vec8_f32 = T.vec(_ELEMENTS_PER_LOAD, f32) + zero_f32 = arith.constant(0.0, type=f32) + + def lds_load(ptr, index): + return fx.ptr_load(ptr + fx.Int64(index)) + + def lds_store(ptr, value, index): + fx.ptr_store(value, ptr + fx.Int64(index)) + + def load_bf16x8(resource, element_index): + dwords = buffer_ops.buffer_load( + resource, + element_index // arith.constant(2, type=i32), + vec_width=4, + dtype=i32, + ) + return vector.bitcast(vec8_bf16, dwords) + + def load_fp8x8_as_bf16(resource, element_index): + packed = ArithValue( + buffer_ops.buffer_load( + resource, + element_index // arith.constant(4, type=i32), + vec_width=2, + dtype=i32, + cache_modifier=weight_cache_modifier, + ) + ) + packed0 = vector.extract(packed, static_position=[0], dynamic_position=[]) + packed1 = vector.extract(packed, static_position=[1], dynamic_position=[]) + weight0_lo = cvt_pk_f32_fp8(res=vec2_f32, src=packed0, word_sel=False) + weight0_hi = cvt_pk_f32_fp8(res=vec2_f32, src=packed0, word_sel=True) + weight1_lo = cvt_pk_f32_fp8(res=vec2_f32, src=packed1, word_sel=False) + weight1_hi = cvt_pk_f32_fp8(res=vec2_f32, src=packed1, word_sel=True) + weight_lo = weight0_lo.shuffle(weight0_hi, [0, 1, 2, 3]) + weight_hi = weight1_lo.shuffle(weight1_hi, [0, 1, 2, 3]) + weight_f32 = weight_lo.shuffle(weight_hi, [0, 1, 2, 3, 4, 5, 6, 7]) + return arith.trunc_f(vec8_bf16, _raw(weight_f32)) + + def wave_reduce_add(value): + reduced = _raw(value) + for offset in (32, 16, 8, 4, 2, 1): + peer = _raw( + ArithValue(reduced).shuffle_xor( + arith.constant(offset, type=i32), + arith.constant(_WAVE_SIZE, type=i32), + ) + ) + reduced = arith.AddFOp(reduced, peer, fastmath=fm_fast).result + return reduced + + def dot_bf16x8(left, right, accumulator): + dot = _raw(accumulator) + for pair_index in range_constexpr(4): + left_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + left, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + left, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + right_pair = vector.from_elements( + vec2_bf16, + [ + vector.extract( + right, + static_position=[pair_index * 2], + dynamic_position=[], + ), + vector.extract( + right, + static_position=[pair_index * 2 + 1], + dynamic_position=[], + ), + ], + ) + dot = llvm.call_intrinsic( + f32, + "llvm.amdgcn.fdot2.f32.bf16", + [ + left_pair, + right_pair, + dot, + arith.constant(False, type=i1), + ], + [], + [], + ) + return ArithValue(dot) + + normalize_thread = arith.cmpi( + CmpIPredicate.ult, + tid, + arith.constant(_NORMALIZE_THREADS, type=i32), + ) + normalize_if = scf.IfOp(normalize_thread, results_=[f32], has_else=True) + with ir.InsertionPoint(normalize_if.then_block): + element_index = tid * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + routed_bf16 = load_bf16x8(routed_rsrc, element_index) + routed_f32 = ArithValue(routed_bf16).extf(vec8_f32) + local_square_sum = (routed_f32 * routed_f32).reduce( + ReductionOp.ADD, fastmath=fm_fast + ) + scf.YieldOp([_raw(local_square_sum)]) + with ir.InsertionPoint(normalize_if.else_block): + scf.YieldOp([zero_f32]) + wave_square_sum = wave_reduce_add(normalize_if.results[0]) + + is_normalize_lane_zero = arith.andi( + arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(0, type=i32), + ), + arith.cmpi( + CmpIPredicate.ult, + wave, + arith.constant(_NORMALIZE_WAVES, type=i32), + ), + ) + lane_zero_if = scf.IfOp(is_normalize_lane_zero) + with ir.InsertionPoint(lane_zero_if.then_block): + lds_store(rms_sums, wave_square_sum, wave) + scf.YieldOp([]) + gpu.barrier() + + is_thread_zero = arith.cmpi(CmpIPredicate.eq, tid, arith.constant(0, type=i32)) + thread_zero_if = scf.IfOp(is_thread_zero) + with ir.InsertionPoint(thread_zero_if.then_block): + total_square_sum = ArithValue(zero_f32) + for wave_index in range_constexpr(_NORMALIZE_WAVES): + total_square_sum = total_square_sum + lds_load( + rms_sums, arith.constant(wave_index, type=i32) + ) + variance = total_square_sum * ArithValue( + arith.constant(1.0 / _LATENT_DIM, type=f32) + ) + inverse = fmath.rsqrt(variance + ArithValue(epsilon), fastmath=fm_fast) + lds_store(inverse_rms, _raw(inverse), arith.constant(0, type=i32)) + scf.YieldOp([]) + gpu.barrier() + + normalize_store_if = scf.IfOp(normalize_thread) + with ir.InsertionPoint(normalize_store_if.then_block): + element_index = tid * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + routed_bf16 = load_bf16x8(routed_rsrc, element_index) + gamma_bf16 = load_bf16x8(rms_weight_rsrc, element_index) + routed_f32 = ArithValue(routed_bf16).extf(vec8_f32) + gamma_f32 = ArithValue(gamma_bf16).extf(vec8_f32) + inverse = ArithValue(lds_load(inverse_rms, arith.constant(0, type=i32))) + normalized = (routed_f32 * gamma_f32 * inverse).truncf(vec8_bf16) + fx.ptr_store(normalized, hidden_lds + element_index) + scf.YieldOp([]) + gpu.barrier() + + first_group = ( + ArithValue(gpu.block_idx.x) * arith.constant(waves_per_block, type=i32) + + wave + ) + for persistent_index in range_constexpr(persistent_iterations): + group = first_group + arith.constant( + persistent_index * groups_per_grid, type=i32 + ) + row_base = group * arith.constant(rows_per_wave, type=i32) + for row_offset in range_constexpr(rows_per_wave): + row = row_base + arith.constant(row_offset, type=i32) + row_in_range = arith.cmpi( + CmpIPredicate.ult, + row, + arith.constant(_HIDDEN_DIM, type=i32), + ) + row_if = scf.IfOp(row_in_range) + with ir.InsertionPoint(row_if.then_block): + local_dot = ArithValue(zero_f32) + for k_iteration in range_constexpr( + _LATENT_DIM // _K_PER_WAVE_ITERATION + ): + k_element = ( + lane + arith.constant(k_iteration * _WAVE_SIZE, type=i32) + ) * arith.constant(_ELEMENTS_PER_LOAD, type=i32) + hidden_bf16 = fx.ptr_load( + hidden_lds + k_element, + result_type=vec8_bf16, + ) + weight_element = ( + row * arith.constant(_LATENT_DIM, type=i32) + k_element + ) + weight_bf16 = load_fp8x8_as_bf16(up_weight_rsrc, weight_element) + local_dot = dot_bf16x8(hidden_bf16, weight_bf16, local_dot) + + scale = buffer_ops.buffer_load( + up_scale_rsrc, row, vec_width=1, dtype=f32 + ) + reduced = ArithValue(wave_reduce_add(local_dot)) * ArithValue(scale) + is_lane_zero = arith.cmpi( + CmpIPredicate.eq, + lane, + arith.constant(0, type=i32), + ) + write_if = scf.IfOp(is_lane_zero) + with ir.InsertionPoint(write_if.then_block): + # Match the BF16 Linear output boundary before the add. + projected_bf16 = arith.trunc_f(T.bf16, _raw(reduced)) + projected_f32 = ArithValue(arith.extf(f32, projected_bf16)) + shared_bf16 = buffer_ops.buffer_load( + shared_rsrc, row, vec_width=1, dtype=T.bf16 + ) + shared_f32 = ArithValue(arith.extf(f32, shared_bf16)) + result = arith.trunc_f(T.bf16, _raw(projected_f32 + shared_f32)) + buffer_ops.buffer_store(result, output_rsrc, row) + scf.YieldOp([]) + scf.YieldOp([]) + + @flyc.jit + def launch_tail( + routed: fx.Pointer, + shared: fx.Pointer, + rms_weight: fx.Pointer, + up_weight: fx.Pointer, + up_scale: fx.Pointer, + output: fx.Pointer, + epsilon: fx.Float32, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + context = CompilationContext.get_current() + if const_expr(waves_per_eu > 0): + for operation in context.gpu_module_body.operations: + if ( + hasattr(operation, "attributes") + and operation.OPERATION_NAME == "gpu.func" + ): + operation.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, int(waves_per_eu) + ) + tail_kernel( + routed, + shared, + rms_weight, + up_weight, + up_scale, + output, + epsilon, + ).launch( + grid=(cu_count, 1, 1), + block=(block_threads, 1, 1), + stream=stream, + ) + + launch_tail.compile_hints = { + "llvm_options": { + "amdgpu-kernarg-preload": AITER_FLYDSL_KERNARG_PRELOAD, + "amdgpu-kernarg-preload-count": AITER_FLYDSL_KERNARG_PRELOAD_COUNT, + }, + } + return launch_tail diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_kda_input_group64.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_kda_input_group64.py new file mode 100644 index 000000000000..ca7410d39056 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_kda_input_group64.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Typed AITER entry point for Kimi-K3 KDA group-64 E4M3 projection.""" + +from __future__ import annotations + +import functools + +import torch +from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + +from .kernels.kimi_k3_kda_input_group64_gfx950 import ( + build_kimi_k3_kda_input_group64_module, +) + +_HIDDEN = 7168 +_PADDED_OUTPUT = 6288 +_LOGICAL_OUTPUT = 6284 +_GROUP = 64 +_GROUPS_PER_ROW = _HIDDEN // _GROUP +_FP8_MAX = 448.0 +_ROWS_PER_WAVE = 2 +_CU_COUNT = 256 +_WEIGHT_CACHE_MODIFIER = 2 + + +def _is_gfx950(device: torch.device) -> bool: + if device.type != "cuda" or not torch.version.hip: + return False + try: + properties = torch.cuda.get_device_properties(device) + except (AssertionError, RuntimeError): + return False + arch = getattr(properties, "gcnArchName", "") + return str(arch).split(":", 1)[0] == "gfx950" + + +def supports_kimi_k3_kda_input_group64( + hidden: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, +) -> bool: + """Fail closed unless every fixed gfx950 contract is satisfied.""" + + tensors = (hidden, weight, scale) + return ( + _is_gfx950(hidden.device) + and hidden.dtype == torch.bfloat16 + and weight.dtype == torch.float8_e4m3fn + and scale.dtype == torch.float32 + and hidden.ndim == 2 + and hidden.shape[0] in (1, 2) + and hidden.shape[1] == _HIDDEN + and tuple(weight.shape) == (_LOGICAL_OUTPUT, _HIDDEN) + and tuple(scale.shape) == (_LOGICAL_OUTPUT, _GROUPS_PER_ROW) + and all(tensor.is_cuda and tensor.is_contiguous() for tensor in tensors) + and len({tensor.device for tensor in tensors}) == 1 + ) + + +def quantize_kimi_k3_kda_input_group64( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Prepack the 6,284 checkpoint-owned rows once, outside decode.""" + + if ( + not weight.is_cuda + or weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (_PADDED_OUTPUT, _HIDDEN) + or not weight.is_contiguous() + ): + raise ValueError("KDA source weight must be contiguous CUDA BF16 [6288,7168]") + padding = weight[_LOGICAL_OUTPUT:] + if bool(torch.count_nonzero(padding).item()): + raise ValueError("KDA projection padding rows must be exactly zero") + source = ( + weight[:_LOGICAL_OUTPUT] + .float() + .reshape(_LOGICAL_OUTPUT, _GROUPS_PER_ROW, _GROUP) + ) + amax = source.abs().amax(dim=-1) + scale = torch.where(amax > 0, amax / _FP8_MAX, torch.ones_like(amax)) + packed = ( + (source / scale[..., None]) + .clamp(min=-_FP8_MAX, max=_FP8_MAX) + .to(torch.float8_e4m3fn) + .reshape(_LOGICAL_OUTPUT, _HIDDEN) + .contiguous() + ) + return packed, scale.contiguous() + + +@functools.lru_cache(maxsize=2) +def _launcher(num_tokens: int): + return build_kimi_k3_kda_input_group64_module( + num_tokens=num_tokens, + rows_per_wave=_ROWS_PER_WAVE, + cu_count=_CU_COUNT, + waves_per_eu=0, + weight_cache_modifier=_WEIGHT_CACHE_MODIFIER, + hidden_to_lds=True, + ) + + +def kimi_k3_kda_input_group64( + hidden: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + output: torch.Tensor | None = None, +) -> torch.Tensor: + """Launch only after the caller has passed the typed support predicate.""" + + if not supports_kimi_k3_kda_input_group64(hidden, weight, scale): + raise ValueError("unsupported Kimi-K3 KDA group64 projection contract") + if output is None: + output = hidden.new_empty((hidden.shape[0], _PADDED_OUTPUT)) + elif ( + output.dtype != torch.bfloat16 + or tuple(output.shape) != (hidden.shape[0], _PADDED_OUTPUT) + or output.device != hidden.device + or not output.is_contiguous() + ): + raise ValueError("output must be contiguous BF16 [M,6288] on the same device") + launcher = _launcher(int(hidden.shape[0])) + launcher( + ptr_arg(hidden), + ptr_arg(weight), + ptr_arg(scale), + ptr_arg(output), + stream=torch.cuda.current_stream(hidden.device), + ) + return output diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_mla_gate.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_mla_gate.py new file mode 100644 index 000000000000..ad676d156343 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_mla_gate.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Kimi-K3 MLA output-gate dispatch.""" + +import functools + +import torch +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +_HIDDEN = 7168 +_OUTPUT = 1536 +_ROWS_PER_BLOCK = 4 +_WAVES_PER_EU = 2 +_WEIGHT_CACHE_MODIFIER = 0 + + +def _is_gfx950_flydsl_available() -> bool: + if not is_flydsl_available(): + return False + try: + return get_gfx_runtime() == "gfx950" + except (AssertionError, KeyError, RuntimeError): + return False + + +def supports_kimi_k3_mla_gate( + hidden: torch.Tensor, + gate_weight: torch.Tensor, + attention_output: torch.Tensor, +) -> bool: + """Return whether the fixed gfx950 BF16 gate specialization owns the call.""" + + tensors = (hidden, gate_weight, attention_output) + return ( + all(tensor.is_cuda for tensor in tensors) + and len({tensor.device for tensor in tensors}) == 1 + and all(tensor.dtype == torch.bfloat16 for tensor in tensors) + and all(tensor.is_contiguous() for tensor in tensors) + and tuple(hidden.shape) == (1, _HIDDEN) + and tuple(gate_weight.shape) == (_OUTPUT, _HIDDEN) + and tuple(attention_output.shape) == (1, _OUTPUT) + and _is_gfx950_flydsl_available() + ) + + +@functools.cache +def _compiled_kimi_k3_mla_gate( + rows_per_block: int, + waves_per_eu: int, + weight_cache_modifier: int, +): + from .kernels.kimi_k3_mla_gate_epilogue_gfx950 import ( + build_kimi_k3_mla_gate_module, + ) + + return build_kimi_k3_mla_gate_module( + rows_per_block, + waves_per_eu, + weight_cache_modifier, + ) + + +def _launch_kimi_k3_mla_gate( + hidden: torch.Tensor, + gate_weight: torch.Tensor, + attention_output: torch.Tensor, + output: torch.Tensor, +) -> torch.Tensor: + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + _compiled_kimi_k3_mla_gate( + _ROWS_PER_BLOCK, + _WAVES_PER_EU, + _WEIGHT_CACHE_MODIFIER, + )( + ptr_arg(hidden), + ptr_arg(gate_weight), + ptr_arg(attention_output), + ptr_arg(output), + stream=torch.cuda.current_stream(hidden.device), + ) + return output + + +def kimi_k3_mla_gate( + hidden: torch.Tensor, + gate_weight: torch.Tensor, + attention_output: torch.Tensor, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Project and apply the BF16 Kimi-K3 full-MLA output gate.""" + + if not supports_kimi_k3_mla_gate(hidden, gate_weight, attention_output): + raise NotImplementedError( + "kimi_k3_mla_gate requires contiguous gfx950 BF16 tensors with " + "shapes (1, 7168), (1536, 7168), and (1, 1536)" + ) + if out is None: + out = torch.empty_like(attention_output) + elif ( + out.device != hidden.device + or out.dtype != torch.bfloat16 + or not out.is_contiguous() + or tuple(out.shape) != (1, _OUTPUT) + ): + raise ValueError( + "out must be contiguous BF16 shape (1, 1536) on the input device" + ) + return _launch_kimi_k3_mla_gate( + hidden, + gate_weight, + attention_output, + out, + ) diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_moe_preroute_fp8.py b/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_moe_preroute_fp8.py new file mode 100644 index 000000000000..a04162c8106f --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/kimi_k3_moe_preroute_fp8.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Kimi-K3 B1 pre-route projections with row-scaled FP8 weights.""" + +from __future__ import annotations + +import functools +import math +import os + +import torch +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +_BATCH_SIZE = 1 +_HIDDEN_SIZE = 7168 +_ROUTED_SIZE = 3584 +_SHARED_GATE_UP_SIZE = 1536 +_SHARED_INTERMEDIATE_SIZE = _SHARED_GATE_UP_SIZE // 2 +_ROUTER_SIZE = 896 + + +def is_kimi_k3_moe_preroute_fp8_available() -> bool: + """Return whether the fixed-shape gfx950 FlyDSL kernels can be built.""" + + return is_flydsl_available() and get_gfx_runtime() == "gfx950" + + +def _same_device(*tensors: torch.Tensor) -> bool: + return len({tensor.device for tensor in tensors}) == 1 + + +def supports_kimi_k3_moe_dual_projection_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, +) -> bool: + """Return whether the fixed Kimi-K3 dual-projection path is supported.""" + + tensors = ( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + ) + return ( + all(tensor.is_cuda for tensor in tensors) + and hidden.dtype == torch.bfloat16 + and routed_weight.dtype == torch.float8_e4m3fn + and routed_scale.dtype == torch.float32 + and shared_weight.dtype == torch.float8_e4m3fn + and shared_scale.dtype == torch.float32 + and tuple(hidden.shape) == (_BATCH_SIZE, _HIDDEN_SIZE) + and tuple(routed_weight.shape) == (_ROUTED_SIZE, _HIDDEN_SIZE) + and tuple(routed_scale.shape) == (_ROUTED_SIZE,) + and tuple(shared_weight.shape) == (_SHARED_GATE_UP_SIZE, _HIDDEN_SIZE) + and tuple(shared_scale.shape) == (_SHARED_GATE_UP_SIZE,) + and all(tensor.is_contiguous() for tensor in tensors) + and _same_device(*tensors) + and is_kimi_k3_moe_preroute_fp8_available() + ) + + +@functools.cache +def _dual_projection_launcher(): + from .kernels.kimi_k3_dual_projection_fp8_gfx950 import ( + build_kimi_k3_b1_dual_projection_fp8_module, + ) + + return build_kimi_k3_b1_dual_projection_fp8_module( + rows_per_wave=2, + cu_count=248, + waves_per_eu=0, + weight_cache_modifier=2, + hidden_to_lds=True, + ) + + +def kimi_k3_moe_dual_projection_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Project one BF16 token with two row-scaled FP8 weight matrices.""" + + if not supports_kimi_k3_moe_dual_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + ): + raise ValueError("unsupported Kimi-K3 dual-projection inputs") + + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + routed_output = hidden.new_empty((_BATCH_SIZE, _ROUTED_SIZE)) + shared_output = hidden.new_empty((_BATCH_SIZE, _SHARED_GATE_UP_SIZE)) + _dual_projection_launcher()( + ptr_arg(hidden), + ptr_arg(routed_weight), + ptr_arg(routed_scale), + ptr_arg(shared_weight), + ptr_arg(shared_scale), + ptr_arg(routed_output), + ptr_arg(shared_output), + stream=torch.cuda.current_stream(hidden.device), + ) + return routed_output, shared_output + + +def _supports_kimi_k3_moe_tri_projection_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, + token_counts: tuple[int, ...], +) -> bool: + tensors = ( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ) + return ( + all(tensor.is_cuda and tensor.is_contiguous() for tensor in tensors) + and hidden.dtype == torch.bfloat16 + and hidden.ndim == 2 + and hidden.shape[0] in token_counts + and hidden.shape[1] == _HIDDEN_SIZE + and routed_weight.dtype == torch.float8_e4m3fn + and tuple(routed_weight.shape) == (_ROUTED_SIZE, _HIDDEN_SIZE) + and routed_scale.dtype == torch.float32 + and tuple(routed_scale.shape) == (_ROUTED_SIZE,) + and shared_weight.dtype == torch.float8_e4m3fn + and tuple(shared_weight.shape) == (_SHARED_GATE_UP_SIZE, _HIDDEN_SIZE) + and shared_scale.dtype == torch.float32 + and tuple(shared_scale.shape) == (_SHARED_GATE_UP_SIZE,) + and router_weight.is_cuda + and router_weight.device == hidden.device + and router_weight.dtype == torch.bfloat16 + and tuple(router_weight.shape) == (_ROUTER_SIZE, _HIDDEN_SIZE) + and router_weight.is_contiguous() + and _same_device(*tensors) + and is_kimi_k3_moe_preroute_fp8_available() + ) + + +def supports_kimi_k3_moe_tri_projection_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, +) -> bool: + """Return whether the production B1/B2 tri-projection is supported.""" + + return _supports_kimi_k3_moe_tri_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + (1,), + ) + + +@functools.cache +def _tri_projection_launcher(num_tokens: int): + from .kernels.kimi_k3_tri_projection_fp8_gfx950 import ( + build_kimi_k3_b1_tri_projection_fp8_module, + ) + + return build_kimi_k3_b1_tri_projection_fp8_module( + num_tokens=num_tokens, + rows_per_wave=1, + cu_count=248, + waves_per_eu=0, + weight_cache_modifier=2, + hidden_to_lds=True, + ) + + +def kimi_k3_moe_tri_projection_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Project routed, shared, and FP32 router outputs in one wide grid.""" + + if not supports_kimi_k3_moe_tri_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ): + raise ValueError("unsupported Kimi-K3 tri-projection inputs") + + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + num_tokens = int(hidden.shape[0]) + routed_output = hidden.new_empty((num_tokens, _ROUTED_SIZE)) + shared_output = hidden.new_empty((num_tokens, _SHARED_GATE_UP_SIZE)) + router_output = hidden.new_empty( + (num_tokens, _ROUTER_SIZE), + dtype=torch.float32, + ) + _tri_projection_launcher(num_tokens)( + ptr_arg(hidden), + ptr_arg(routed_weight), + ptr_arg(routed_scale), + ptr_arg(shared_weight), + ptr_arg(shared_scale), + ptr_arg(router_weight), + ptr_arg(routed_output), + ptr_arg(shared_output), + ptr_arg(router_output), + stream=torch.cuda.current_stream(hidden.device), + ) + return routed_output, shared_output, router_output + + +def supports_kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, +) -> bool: + return _supports_kimi_k3_moe_tri_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + (2, 4), + ) + + +@functools.cache +def _tri_cooperative_preactivated_launcher( + num_tokens: int, + situ_beta: float, + situ_linear_beta: float, + fast_situ: bool, +): + from .kernels.kimi_k3_tri_projection_multitoken_gfx950 import ( + build_kimi_k3_multitoken_tri_projection_module, + ) + + return build_kimi_k3_multitoken_tri_projection_module( + num_tokens=num_tokens, + token_tile=num_tokens, + cu_count=int(os.environ.get("SGLANG_K3_PREROUTE_COOP_CU", "256")), + waves_per_block=int(os.environ.get("SGLANG_K3_PREROUTE_COOP_WPB", "8")), + waves_per_eu=int(os.environ.get("SGLANG_K3_PREROUTE_COOP_WPE", "3")), + weight_cache_modifier=int(os.environ.get("SGLANG_K3_PREROUTE_COOP_WCM", "3")), + interleaved_shared_pairs=True, + fast_situ=fast_situ, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + +def kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, + *, + situ_beta: float, + situ_linear_beta: float, + fast_situ: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if ( + not math.isfinite(situ_beta) + or not math.isfinite(situ_linear_beta) + or situ_beta <= 0.0 + or situ_linear_beta <= 0.0 + ): + raise ValueError("SiTU beta values must be finite and positive") + if not supports_kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ): + raise ValueError( + "unsupported Kimi-K3 cooperative preactivated tri-projection inputs" + ) + + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + num_tokens = int(hidden.shape[0]) + routed_output = hidden.new_empty((num_tokens, _ROUTED_SIZE)) + shared_output = hidden.new_empty((num_tokens, _SHARED_INTERMEDIATE_SIZE)) + router_output = hidden.new_empty((num_tokens, _ROUTER_SIZE), dtype=torch.float32) + _tri_cooperative_preactivated_launcher( + num_tokens, + float(situ_beta), + float(situ_linear_beta), + bool(fast_situ), + )( + ptr_arg(hidden), + ptr_arg(routed_weight), + ptr_arg(routed_scale), + ptr_arg(shared_weight), + ptr_arg(shared_scale), + ptr_arg(router_weight), + ptr_arg(routed_output), + ptr_arg(shared_output), + ptr_arg(router_output), + stream=torch.cuda.current_stream(hidden.device), + ) + return routed_output, shared_output, router_output + + +def supports_kimi_k3_shared_down_fp8( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, +) -> bool: + """Return whether the fused Kimi-K3 SiTU/shared-down path is supported.""" + + tensors = (gate_up, weight, weight_scale) + return ( + gate_up.is_cuda + and gate_up.dtype == torch.bfloat16 + and gate_up.ndim == 2 + and gate_up.shape[0] == 1 + and gate_up.shape[1] == _SHARED_GATE_UP_SIZE + and gate_up.is_contiguous() + and supports_kimi_k3_shared_down_fp8_weight( + weight, + weight_scale, + device=gate_up.device, + ) + and _same_device(*tensors) + ) + + +def supports_kimi_k3_shared_down_fp8_weight( + weight: torch.Tensor, + weight_scale: torch.Tensor, + *, + device: torch.device | None = None, +) -> bool: + """Return whether a shared-down weight pair has the kernel's contract.""" + + tensors = (weight, weight_scale) + return ( + all(tensor.is_cuda for tensor in tensors) + and weight.dtype == torch.float8_e4m3fn + and weight_scale.dtype == torch.float32 + and tuple(weight.shape) == (_HIDDEN_SIZE, _SHARED_INTERMEDIATE_SIZE) + and tuple(weight_scale.shape) == (_HIDDEN_SIZE,) + and all(tensor.is_contiguous() for tensor in tensors) + and _same_device(*tensors) + and (device is None or weight.device == device) + and is_kimi_k3_moe_preroute_fp8_available() + ) + + +@functools.cache +def _shared_down_launcher( + num_tokens: int, + situ_beta: float, + situ_linear_beta: float, +): + from .kernels.kimi_k3_shared_down_fp8_gfx950 import ( + build_kimi_k3_b1_shared_down_fp8_module, + ) + + return build_kimi_k3_b1_shared_down_fp8_module( + num_tokens=num_tokens, + rows_per_wave=1, + cu_count=248, + waves_per_eu=0, + weight_cache_modifier=2, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + +def kimi_k3_shared_down_fp8( + gate_up: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + *, + situ_beta: float, + situ_linear_beta: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Apply SiTU and project with a row-scaled FP8 weight matrix.""" + + if ( + not math.isfinite(situ_beta) + or not math.isfinite(situ_linear_beta) + or situ_beta <= 0.0 + or situ_linear_beta <= 0.0 + ): + raise ValueError("SiTU beta values must be finite and positive") + if not supports_kimi_k3_shared_down_fp8( + gate_up, + weight, + weight_scale, + ): + raise ValueError("unsupported Kimi-K3 shared-down inputs") + + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + if out is None: + output = gate_up.new_empty((gate_up.shape[0], _HIDDEN_SIZE)) + elif ( + out.device != gate_up.device + or out.dtype != torch.bfloat16 + or tuple(out.shape) != (gate_up.shape[0], _HIDDEN_SIZE) + or not out.is_contiguous() + ): + raise ValueError("out must be contiguous BF16 [M,7168] on the same device") + else: + output = out + _shared_down_launcher( + int(gate_up.shape[0]), + float(situ_beta), + float(situ_linear_beta), + )( + ptr_arg(gate_up), + ptr_arg(weight), + ptr_arg(weight_scale), + ptr_arg(output), + stream=torch.cuda.current_stream(gate_up.device), + ) + return output diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/latent_moe_tail_fp8.py b/python/sglang/kernels/ops/kimi_k3/flydsl/latent_moe_tail_fp8.py new file mode 100644 index 000000000000..67e04c8de373 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/latent_moe_tail_fp8.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Typed row-scaled-FP8 Kimi-K3 latent-tail entry point for gfx950.""" + +from __future__ import annotations + +import functools +import math + +import torch +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +_LATENT_DIM = 3584 +_HIDDEN_DIM = 7168 +_FP8_MAX = 448.0 + + +def quantize_latent_moe_tail_weight( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack one contiguous BF16 up-projection to row-scaled OCP E4M3.""" + + if ( + not weight.is_cuda + or weight.dtype != torch.bfloat16 + or tuple(weight.shape) != (_HIDDEN_DIM, _LATENT_DIM) + or not weight.is_contiguous() + ): + raise ValueError( + "latent-tail source weight must be contiguous CUDA BF16 [7168,3584]" + ) + weight_f32 = weight.float() + amax = weight_f32.abs().amax(dim=1) + scale = torch.where( + amax > 0, + amax / _FP8_MAX, + torch.ones_like(amax), + ) + packed = ( + (weight_f32 / scale[:, None]) + .clamp(min=-_FP8_MAX, max=_FP8_MAX) + .to(torch.float8_e4m3fn) + .contiguous() + ) + return packed, scale.contiguous() + + +def supports_latent_moe_tail_fp8( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + up_scale: torch.Tensor, + epsilon: float, +) -> bool: + """Fail closed unless the exact MI355X TP8 B1 contract is present.""" + + tensors = (routed, shared, rms_weight, up_weight, up_scale) + return ( + all(tensor.is_cuda for tensor in tensors) + and len({tensor.device for tensor in tensors}) == 1 + and all(tensor.is_contiguous() for tensor in tensors) + and routed.dtype == torch.bfloat16 + and shared.dtype == torch.bfloat16 + and rms_weight.dtype == torch.bfloat16 + and up_weight.dtype == torch.float8_e4m3fn + and up_scale.dtype == torch.float32 + and tuple(routed.shape) == (1, _LATENT_DIM) + and tuple(shared.shape) == (1, _HIDDEN_DIM) + and tuple(rms_weight.shape) == (_LATENT_DIM,) + and tuple(up_weight.shape) == (_HIDDEN_DIM, _LATENT_DIM) + and tuple(up_scale.shape) == (_HIDDEN_DIM,) + and math.isfinite(epsilon) + and epsilon > 0.0 + and is_flydsl_available() + and get_gfx_runtime() == "gfx950" + ) + + +@functools.cache +def _compiled_latent_moe_tail_fp8(): + from .kernels.latent_moe_tail_fp8_gfx950 import ( + build_b1_latent_moe_tail_fp8_persistent_module, + ) + + return build_b1_latent_moe_tail_fp8_persistent_module( + rows_per_wave=1, + cu_count=240, + waves_per_eu=2, + weight_cache_modifier=2, + ) + + +def latent_moe_tail_fp8( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + up_scale: torch.Tensor, + epsilon: float, + *, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Fuse RMSNorm, FP8-weight GEMV, BF16 materialization, and shared add.""" + + if not supports_latent_moe_tail_fp8( + routed, + shared, + rms_weight, + up_weight, + up_scale, + epsilon, + ): + raise NotImplementedError("unsupported Kimi-K3 FP8 latent-tail contract") + if out is None: + out = torch.empty_like(shared) + elif ( + out.device != routed.device + or out.dtype != torch.bfloat16 + or not out.is_contiguous() + or tuple(out.shape) != (1, _HIDDEN_DIM) + ): + raise ValueError( + "out must be contiguous BF16 shape (1,7168) on the input device" + ) + + from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg + + _compiled_latent_moe_tail_fp8()( + ptr_arg(routed), + ptr_arg(shared), + ptr_arg(rms_weight), + ptr_arg(up_weight), + ptr_arg(up_scale), + ptr_arg(out), + float(epsilon), + stream=torch.cuda.current_stream(routed.device), + ) + return out + + +__all__ = [ + "latent_moe_tail_fp8", + "quantize_latent_moe_tail_weight", + "supports_latent_moe_tail_fp8", +] diff --git a/python/sglang/kernels/ops/kimi_k3/flydsl/source.py b/python/sglang/kernels/ops/kimi_k3/flydsl/source.py new file mode 100644 index 000000000000..7337acf57f38 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/flydsl/source.py @@ -0,0 +1,27 @@ +"""Select SGLang-vendored or upstream AITER Kimi-K3 FlyDSL operators.""" + +import importlib +import os + + +def load_module(local_module: str, aiter_module: str): + mode = os.environ.get("SGLANG_K3_FLYDSL_SOURCE", "auto").lower() + if mode not in ("auto", "sglang", "aiter"): + raise ValueError( + "SGLANG_K3_FLYDSL_SOURCE must be one of auto, sglang, or aiter" + ) + + candidates = ( + ((local_module, "sglang"), (aiter_module, "aiter")) + if mode in ("auto", "sglang") + else ((aiter_module, "aiter"),) + ) + errors = [] + for module_name, source in candidates: + if mode == "sglang" and source != "sglang": + continue + try: + return importlib.import_module(module_name) + except (ImportError, ModuleNotFoundError) as error: + errors.append(f"{source}: {error}") + raise ImportError("Kimi-K3 FlyDSL source unavailable: " + "; ".join(errors)) diff --git a/python/sglang/kernels/ops/kimi_k3/kda_group64_aiter_hip.py b/python/sglang/kernels/ops/kimi_k3/kda_group64_aiter_hip.py new file mode 100644 index 000000000000..f5adc569bf06 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/kda_group64_aiter_hip.py @@ -0,0 +1,82 @@ +"""Fail-closed adapter for AITER's gfx950 Kimi-K3 KDA input projection.""" + +from __future__ import annotations + +import os + +import torch + +from sglang.srt.utils import is_hip + + +def enabled() -> bool: + return is_hip() and os.environ.get("SGLANG_K3_AITER_KDA_GROUP64", "0").lower() in ( + "1", + "true", + ) + + +def _b2_enabled() -> bool: + return os.environ.get("SGLANG_K3_AITER_B2_FUSIONS", "0").lower() in ( + "1", + "true", + ) + + +def _ops(): + try: + from sglang.kernels.ops.kimi_k3.flydsl.source import load_module + + module = load_module( + "sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_kda_input_group64", + "aiter.ops.flydsl.kimi_k3_kda_input_group64", + ) + except (ImportError, ModuleNotFoundError): + return None, None, None + return ( + module.kimi_k3_kda_input_group64, + module.quantize_kimi_k3_kda_input_group64, + module.supports_kimi_k3_kda_input_group64, + ) + + +def pack(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + _, quantize, _ = _ops() + if quantize is None: + raise RuntimeError("AITER KDA group64 quantizer is unavailable") + return quantize(weight) + + +def covered(hidden: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> bool: + if not enabled(): + return False + if hidden.shape[0] > 1 and (hidden.shape[0] != 2 or not _b2_enabled()): + return False + _, _, supports = _ops() + return bool(supports is not None and supports(hidden, weight, scale)) + + +def run( + hidden: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + *, + output: torch.Tensor | None = None, +) -> torch.Tensor: + op, _, _ = _ops() + if op is None: + raise RuntimeError("AITER KDA group64 projection is unavailable") + return op(hidden, weight, scale, output=output) + + +def warmup(weight: torch.Tensor, scale: torch.Tensor) -> None: + if not enabled(): + return + token_buckets = (1, 2) if _b2_enabled() else (1,) + for num_tokens in token_buckets: + hidden = torch.zeros( + (num_tokens, 7168), dtype=torch.bfloat16, device=weight.device + ) + if covered(hidden, weight, scale): + run(hidden, weight, scale) + torch.cuda.synchronize(weight.device) diff --git a/python/sglang/kernels/ops/kimi_k3/latent_tail_aiter_hip.py b/python/sglang/kernels/ops/kimi_k3/latent_tail_aiter_hip.py new file mode 100644 index 000000000000..05ef84ac00b7 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/latent_tail_aiter_hip.py @@ -0,0 +1,99 @@ +"""Fail-closed adapter for AITER's gfx950 Kimi-K3 FP8 latent tail.""" + +from __future__ import annotations + +import os + +import torch + +from sglang.srt.utils import is_hip + + +def enabled() -> bool: + return is_hip() and os.environ.get( + "SGLANG_K3_AITER_LATENT_TAIL_FP8", "0" + ).lower() in ("1", "true") + + +def _ops(): + try: + from sglang.kernels.ops.kimi_k3.flydsl.source import load_module + + module = load_module( + "sglang.kernels.ops.kimi_k3.flydsl.latent_moe_tail_fp8", + "aiter.ops.flydsl.latent_moe_tail_fp8", + ) + except (ImportError, ModuleNotFoundError): + return None, None, None + return ( + module.latent_moe_tail_fp8, + module.quantize_latent_moe_tail_weight, + module.supports_latent_moe_tail_fp8, + ) + + +def pack(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + _, quantize, _ = _ops() + if quantize is None: + raise RuntimeError("AITER latent-tail quantizer is unavailable") + return quantize(weight.contiguous()) + + +def covered( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + up_scale: torch.Tensor, + epsilon: float, +) -> bool: + if not enabled(): + return False + _, _, supports = _ops() + return bool( + supports is not None + and supports( + routed, + shared, + rms_weight, + up_weight, + up_scale, + epsilon, + ) + ) + + +def run( + routed: torch.Tensor, + shared: torch.Tensor, + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + up_scale: torch.Tensor, + epsilon: float, +) -> torch.Tensor: + op, _, _ = _ops() + if op is None: + raise RuntimeError("AITER latent-tail fusion is unavailable") + return op( + routed, + shared, + rms_weight, + up_weight, + up_scale, + epsilon, + ) + + +def warmup( + rms_weight: torch.Tensor, + up_weight: torch.Tensor, + up_scale: torch.Tensor, + epsilon: float, +) -> None: + if not enabled(): + return + routed = torch.zeros((1, 3584), dtype=torch.bfloat16, device=up_weight.device) + shared = torch.zeros((1, 7168), dtype=torch.bfloat16, device=up_weight.device) + if covered(routed, shared, rms_weight, up_weight, up_scale, epsilon): + run(routed, shared, rms_weight, up_weight, up_scale, epsilon) + torch.cuda.synchronize(up_weight.device) diff --git a/python/sglang/kernels/ops/kimi_k3/mla_gate_aiter_hip.py b/python/sglang/kernels/ops/kimi_k3/mla_gate_aiter_hip.py new file mode 100644 index 000000000000..7951a2fda7c3 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/mla_gate_aiter_hip.py @@ -0,0 +1,63 @@ +"""Fail-closed adapter for AITER's gfx950 Kimi-K3 MLA output gate.""" + +from __future__ import annotations + +import os + +import torch + +from sglang.srt.utils import is_hip + + +def enabled() -> bool: + return is_hip() and os.environ.get("SGLANG_K3_AITER_MLA_GATE", "0").lower() in ( + "1", + "true", + ) + + +def _ops(): + try: + from sglang.kernels.ops.kimi_k3.flydsl.source import load_module + + module = load_module( + "sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_mla_gate", + "aiter.ops.flydsl.kimi_k3_mla_gate", + ) + except (ImportError, ModuleNotFoundError): + return None, None + return module.kimi_k3_mla_gate, module.supports_kimi_k3_mla_gate + + +def covered( + hidden: torch.Tensor, + gate_weight: torch.Tensor, + attention_output: torch.Tensor, +) -> bool: + if not enabled(): + return False + _, supports = _ops() + return bool( + supports is not None and supports(hidden, gate_weight, attention_output) + ) + + +def run( + hidden: torch.Tensor, + gate_weight: torch.Tensor, + attention_output: torch.Tensor, +) -> torch.Tensor: + op, _ = _ops() + if op is None: + raise RuntimeError("AITER Kimi-K3 MLA gate is unavailable") + return op(hidden, gate_weight, attention_output, out=attention_output) + + +def warmup(gate_weight: torch.Tensor) -> None: + if not enabled() or tuple(gate_weight.shape) != (1536, 7168): + return + hidden = torch.zeros((1, 7168), dtype=torch.bfloat16, device=gate_weight.device) + attention = torch.zeros((1, 1536), dtype=torch.bfloat16, device=gate_weight.device) + if covered(hidden, gate_weight, attention): + run(hidden, gate_weight, attention) + torch.cuda.synchronize(gate_weight.device) diff --git a/python/sglang/kernels/ops/kimi_k3/moe_preroute_aiter_hip.py b/python/sglang/kernels/ops/kimi_k3/moe_preroute_aiter_hip.py new file mode 100644 index 000000000000..04d37ea343d7 --- /dev/null +++ b/python/sglang/kernels/ops/kimi_k3/moe_preroute_aiter_hip.py @@ -0,0 +1,280 @@ +"""Fail-closed adapter for AITER's gfx950 Kimi-K3 FP8 MoE front.""" + +from __future__ import annotations + +import os + +import torch + +from sglang.srt.utils import is_hip + + +def enabled() -> bool: + return is_hip() and os.environ.get( + "SGLANG_K3_AITER_MOE_PREROUTE_FP8", "0" + ).lower() in ("1", "true") + + +def cooperative_preactivated_enabled() -> bool: + return ( + enabled() + and os.environ.get("SGLANG_K3_FLYDSL_SOURCE", "auto").lower() != "aiter" + and os.environ.get("SGLANG_K3_PREROUTE_PREACTIVATED_SHARED", "0").lower() + in ("1", "true") + ) + + +def _ops(): + try: + from sglang.kernels.ops.kimi_k3.flydsl.source import load_module + + module = load_module( + "sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_moe_preroute_fp8", + "aiter.ops.flydsl.kimi_k3_moe_preroute_fp8", + ) + except (ImportError, ModuleNotFoundError): + return None, None, None, None + return ( + module.kimi_k3_moe_tri_projection_fp8, + module.kimi_k3_shared_down_fp8, + module.supports_kimi_k3_moe_tri_projection_fp8, + module.supports_kimi_k3_shared_down_fp8, + ) + + +def _preactivated_ops(): + try: + from sglang.kernels.ops.kimi_k3.flydsl.source import load_module + + module = load_module( + "sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_moe_preroute_fp8", + "aiter.ops.flydsl.kimi_k3_moe_preroute_fp8", + ) + except (ImportError, ModuleNotFoundError): + return None, None + return ( + getattr( + module, + "kimi_k3_moe_tri_projection_cooperative_preactivated_fp8", + None, + ), + getattr( + module, + "supports_kimi_k3_moe_tri_projection_cooperative_preactivated_fp8", + None, + ), + ) + + +def cooperative_preactivated_tri_covered( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, +) -> bool: + if not cooperative_preactivated_enabled() or hidden.shape[0] not in (2, 4): + return False + _, supports = _preactivated_ops() + return bool( + supports is not None + and supports( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ) + ) + + +def run_tri_cooperative_preactivated( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, + *, + situ_beta: float, + situ_linear_beta: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + op, _ = _preactivated_ops() + if op is None: + raise RuntimeError("Kimi-K3 cooperative preactivated projection is unavailable") + return op( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + fast_situ=True, + ) + + +def tri_covered( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, +) -> bool: + if not enabled(): + return False + num_tokens = hidden.shape[0] + if num_tokens != 1: + return False + _, _, supports, _ = _ops() + return bool( + supports is not None + and supports( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ) + ) + + +def run_tri( + hidden: torch.Tensor, + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + op, _, _, _ = _ops() + if op is None: + raise RuntimeError("AITER Kimi-K3 pre-route projection is unavailable") + return op( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ) + + +def shared_down_covered( + gate_up: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, +) -> bool: + if not enabled(): + return False + if gate_up.shape[0] != 1: + return False + _, _, _, supports = _ops() + return bool(supports is not None and supports(gate_up, weight, scale)) + + +def run_shared_down( + gate_up: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + *, + situ_beta: float, + situ_linear_beta: float, + out: torch.Tensor, +) -> torch.Tensor: + _, op, _, _ = _ops() + if op is None: + raise RuntimeError("AITER Kimi-K3 shared-down is unavailable") + return op( + gate_up, + weight, + scale, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + out=out, + ) + + +def warmup( + routed_weight: torch.Tensor, + routed_scale: torch.Tensor, + shared_weight: torch.Tensor, + shared_scale: torch.Tensor, + router_weight: torch.Tensor, + shared_down_weight: torch.Tensor, + shared_down_scale: torch.Tensor, + shared_interleaved_weight: torch.Tensor | None = None, + shared_interleaved_scale: torch.Tensor | None = None, + *, + situ_beta: float, + situ_linear_beta: float, +) -> None: + if not enabled(): + return + token_buckets = [1] + if cooperative_preactivated_enabled(): + token_buckets.append(2) + token_buckets.append(4) + for num_tokens in token_buckets: + hidden = torch.zeros( + (num_tokens, 7168), + dtype=torch.bfloat16, + device=routed_weight.device, + ) + if ( + num_tokens in (2, 4) + and shared_interleaved_weight is not None + and shared_interleaved_scale is not None + and cooperative_preactivated_tri_covered( + hidden, + routed_weight, + routed_scale, + shared_interleaved_weight, + shared_interleaved_scale, + router_weight, + ) + ): + run_tri_cooperative_preactivated( + hidden, + routed_weight, + routed_scale, + shared_interleaved_weight, + shared_interleaved_scale, + router_weight, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + continue + if not tri_covered( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ): + continue + _, gate_up, _ = run_tri( + hidden, + routed_weight, + routed_scale, + shared_weight, + shared_scale, + router_weight, + ) + out = hidden.new_empty((num_tokens, 7168)) + if shared_down_covered(gate_up, shared_down_weight, shared_down_scale): + run_shared_down( + gate_up, + shared_down_weight, + shared_down_scale, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + out=out, + ) + torch.cuda.synchronize(hidden.device) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 3943a58124b0..a44f01bd3441 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -295,11 +295,11 @@ def _pool_is_vec5d(pool): self.forward_metadata: ForwardMetadata = None if self.use_mla: - _valid_heads = self.num_head in (4, 8) or ( + _valid_heads = 1 <= self.num_head < 16 or ( self.num_head % 16 == 0 and 16 <= self.num_head <= 128 ) assert _valid_heads, ( - f"Aiter MLA supports num_head of 4, 8, or multiples of 16 " + f"Aiter MLA supports num_head below 16 or multiples of 16 " f"in [16, 128].\n" f"Provided {self.num_head} number of heads.\n" "Try adjusting tensor_parallel_size value." @@ -760,22 +760,28 @@ def _mla_decode_fwd_with_head_pad( ): """Wrap mla_decode_fwd with head-dimension padding for num_head < 16. - When head_repeat_factor > 1 (i.e. num_head is 4 or 8), q is - repeat-interleaved to reach num_head_padded (16) before the kernel - call, and the corresponding output columns are sliced back afterward. + Head counts that divide 16 (e.g. 4 or 8) retain the historical + repeat-interleave path. Other counts (e.g. Kimi-K3 TP8's 12 heads) are + zero-padded to 16 and sliced back after the kernel call. q / o must already be shaped (..., num_head, head_dim). """ - if self.head_repeat_factor > 1: - q_in = q.repeat_interleave(self.head_repeat_factor, dim=1) + n_heads = layer.tp_q_head_num + if n_heads < 16: + if 16 % n_heads == 0: + q_in = q.repeat_interleave(16 // n_heads, dim=1) + select = slice(None, None, 16 // n_heads) + else: + q_in = torch.nn.functional.pad(q, (0, 0, 0, 16 - n_heads)) + select = slice(0, n_heads) o = q.new_empty( - (q.shape[0], self.num_head_padded, layer.v_head_dim), + (q.shape[0], 16, layer.v_head_dim), dtype=self.input_dtype, ) mla_decode_fwd(q_in, k_buffer_flat, o, **kwargs) - return o[:, :: self.head_repeat_factor, :] + return o[:, select, :] else: o = q.new_empty( - (q.shape[0], layer.tp_q_head_num, layer.v_head_dim), + (q.shape[0], n_heads, layer.v_head_dim), dtype=self.input_dtype, ) mla_decode_fwd(q, k_buffer_flat, o, **kwargs) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 378f3cbf65a5..cf27e69dc9d3 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -878,4 +878,10 @@ def _skip_rope_for_aiter_fused_mla(self: DeepseekV2AttentionMLA) -> bool: when running aiter-backend MLA on gfx95 (i.e., the `else` branch in forward_absorb_rocm_core that calls fused_qk_rope_cat_and_cache_mla). """ - return _use_aiter_gfx95 and self.current_attention_backend == "aiter" + # NoPE models (rotary_emb=None, e.g. Kimi-K3) have no rope for the + # fused kernel to apply and no cos/sin cache to pass to it. + return ( + self.rotary_emb is not None + and _use_aiter_gfx95 + and self.current_attention_backend == "aiter" + ) diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index f910195d08f2..b8b81fc99d04 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -134,9 +134,15 @@ logger = logging.getLogger(__name__) _is_hip = is_hip() _is_npu = is_npu() -_aiter_k3_opt = get_bool_env_var("SGLANG_AITER_K3_OPT") +_aiter_k3_opt = _is_hip and get_bool_env_var("SGLANG_AITER_K3_OPT") _k3_shared_experts_attn_tp = envs.SGLANG_K3_SHARED_EXPERTS_ATTN_TP.get() _k3_dense_mlp_attn_tp = envs.SGLANG_K3_DENSE_MLP_ATTN_TP.get() +_aiter_mla_gate = _is_hip and get_bool_env_var("SGLANG_K3_AITER_MLA_GATE") +_aiter_kda_group64 = _is_hip and get_bool_env_var("SGLANG_K3_AITER_KDA_GROUP64") +_aiter_moe_preroute_fp8 = _is_hip and get_bool_env_var( + "SGLANG_K3_AITER_MOE_PREROUTE_FP8" +) +_aiter_latent_tail_fp8 = _is_hip and get_bool_env_var("SGLANG_K3_AITER_LATENT_TAIL_FP8") def _cdiv(a: int, b: int) -> int: @@ -586,6 +592,18 @@ def __init__( ) else: self.shared_experts = None + self._preroute_routed_weight = None + self._preroute_routed_scale = None + self._preroute_shared_weight = None + self._preroute_shared_scale = None + self._preroute_shared_interleaved_weight = None + self._preroute_shared_interleaved_scale = None + self._preroute_shared_down_weight = None + self._preroute_shared_down_scale = None + self._latent_tail_weight = None + self._latent_tail_scale = None + self._situ_beta = float(config.activation_situ_beta) + self._situ_linear_beta = float(config.activation_situ_linear_beta) # SBO (single batch overlap): the shared experts read a fixed slab of # weights the routed path never touches (bf16 — the checkpoint leaves @@ -718,6 +736,87 @@ def _merge_front_weights(self) -> None: ): self.__dict__.pop(prop, None) + def _prepare_preroute_fp8(self) -> None: + if ( + not _aiter_moe_preroute_fp8 + or not self.use_latent_moe + or self.shared_experts is None + ): + return + from sglang.kernels.ops.kimi_k3 import ( + moe_preroute_aiter_hip, + ) + from sglang.kernels.ops.kimi_k3.aiter_fusion import ( + quantize_fp8_rows, + ) + + routed = self.routed_expert_down_proj.weight + shared = self.shared_experts.gate_up_proj.weight + shared_down = self.shared_experts.down_proj.weight + if ( + tuple(routed.shape) != (3584, 7168) + or tuple(shared.shape) != (1536, 7168) + or tuple(shared_down.shape) != (7168, 768) + or tuple(self.gate.weight.shape) != (896, 7168) + ): + return + self._preroute_routed_weight, self._preroute_routed_scale = quantize_fp8_rows( + routed.contiguous() + ) + self._preroute_shared_weight, self._preroute_shared_scale = quantize_fp8_rows( + shared.contiguous() + ) + if moe_preroute_aiter_hip.cooperative_preactivated_enabled(): + self._preroute_shared_interleaved_weight = ( + self._preroute_shared_weight.view(2, 768, 7168) + .permute(1, 0, 2) + .contiguous() + .view(1536, 7168) + ) + self._preroute_shared_interleaved_scale = ( + self._preroute_shared_scale.view(2, 768).t().contiguous().view(1536) + ) + ( + self._preroute_shared_down_weight, + self._preroute_shared_down_scale, + ) = quantize_fp8_rows(shared_down.contiguous()) + moe_preroute_aiter_hip.warmup( + self._preroute_routed_weight, + self._preroute_routed_scale, + self._preroute_shared_weight, + self._preroute_shared_scale, + self.gate.weight, + self._preroute_shared_down_weight, + self._preroute_shared_down_scale, + self._preroute_shared_interleaved_weight, + self._preroute_shared_interleaved_scale, + situ_beta=self._situ_beta, + situ_linear_beta=self._situ_linear_beta, + ) + + def _prepare_latent_tail_fp8(self) -> None: + if ( + not _aiter_latent_tail_fp8 + or not self.fuse_ar_norm + or self.routed_expert_up_proj is None + or self.routed_expert_norm is None + ): + return + from sglang.kernels.ops.kimi_k3 import latent_tail_aiter_hip + + if tuple(self.routed_expert_up_proj.weight.shape) != (7168, 3584): + return + self._latent_tail_weight, self._latent_tail_scale = latent_tail_aiter_hip.pack( + self.routed_expert_up_proj.weight + ) + norm_weight, epsilon = self._get_fused_norm_params() + latent_tail_aiter_hip.warmup( + norm_weight, + self._latent_tail_weight, + self._latent_tail_scale, + epsilon, + ) + @cached_property def _routed_needs_reduce(self): return self.tp_size > 1 and get_moe_a2a_backend().is_none() @@ -1166,13 +1265,48 @@ def _forward_routed_deferred(self, hidden_states, router_logits, routed_input): finally: route_quant_handoff.clear() - def _forward_shared(self, gate_up, shared_output): + def _forward_shared( + self, + gate_up, + shared_output, + *, + preactivated: bool = False, + ): shared = self.shared_experts if TYPE_CHECKING: assert shared is not None and isinstance( shared.down_proj.weight, torch.Tensor ) assert shared is not None + if preactivated: + _k3_bf16_gemm( + gate_up, + shared.down_proj.weight, + out=shared_output, + ) + return + if ( + self._preroute_shared_down_weight is not None + and self._preroute_shared_down_scale is not None + ): + from sglang.kernels.ops.kimi_k3 import ( + moe_preroute_aiter_hip, + ) + + if moe_preroute_aiter_hip.shared_down_covered( + gate_up, + self._preroute_shared_down_weight, + self._preroute_shared_down_scale, + ): + moe_preroute_aiter_hip.run_shared_down( + gate_up, + self._preroute_shared_down_weight, + self._preroute_shared_down_scale, + situ_beta=self._situ_beta, + situ_linear_beta=self._situ_linear_beta, + out=shared_output, + ) + return _k3_bf16_gemm( shared.act_fn(gate_up), shared.down_proj.weight, @@ -1205,14 +1339,71 @@ def _forward_fused( ) num_tokens, hidden_size = hidden_states.shape - fused = _k3_bf16_gemm( - hidden_states, - self._front_w, - out_dtype=torch.float32 if self._front_fp32 else None, - ) - gate_up, router_logits, routed_input = torch.split( - fused, self._front_sizes, dim=-1 - ) + preroute = None + shared_is_preactivated = False + if ( + num_tokens <= 4 + and self._preroute_routed_weight is not None + and self._preroute_routed_scale is not None + and self._preroute_shared_weight is not None + and self._preroute_shared_scale is not None + ): + from sglang.kernels.ops.kimi_k3 import ( + moe_preroute_aiter_hip, + ) + + if ( + self._preroute_shared_interleaved_weight is not None + and self._preroute_shared_interleaved_scale is not None + and moe_preroute_aiter_hip.cooperative_preactivated_tri_covered( + hidden_states, + self._preroute_routed_weight, + self._preroute_routed_scale, + self._preroute_shared_interleaved_weight, + self._preroute_shared_interleaved_scale, + self.gate.weight, + ) + ): + routed_input, gate_up, router_logits = ( + moe_preroute_aiter_hip.run_tri_cooperative_preactivated( + hidden_states, + self._preroute_routed_weight, + self._preroute_routed_scale, + self._preroute_shared_interleaved_weight, + self._preroute_shared_interleaved_scale, + self.gate.weight, + situ_beta=self._situ_beta, + situ_linear_beta=self._situ_linear_beta, + ) + ) + shared_is_preactivated = True + preroute = True + elif moe_preroute_aiter_hip.tri_covered( + hidden_states, + self._preroute_routed_weight, + self._preroute_routed_scale, + self._preroute_shared_weight, + self._preroute_shared_scale, + self.gate.weight, + ): + routed_input, gate_up, router_logits = moe_preroute_aiter_hip.run_tri( + hidden_states, + self._preroute_routed_weight, + self._preroute_routed_scale, + self._preroute_shared_weight, + self._preroute_shared_scale, + self.gate.weight, + ) + preroute = True + if preroute is None: + fused = _k3_bf16_gemm( + hidden_states, + self._front_w, + out_dtype=torch.float32 if self._front_fp32 else None, + ) + gate_up, router_logits, routed_input = torch.split( + fused, self._front_sizes, dim=-1 + ) if num_tokens > 1 and _is_hip and not _aiter_k3_opt: router_logits = router_logits.contiguous() if self._moe_front_needs_dense_bf16: @@ -1237,12 +1428,18 @@ def _forward_fused( latent = buf[:latent_numel].view(num_tokens, self.moe_hidden_size) shared_output = buf[latent_numel:].view(num_tokens, hidden_size) + use_latent_tail = ( + num_tokens == 1 + and self._latent_tail_weight is not None + and self._latent_tail_scale is not None + ) fused_norm = False if self.alt_stream is not None and k3_ar_fusion.enabled(): defer_finalize = ( self._defer_moe_finalize and self.fuse_ar_norm and k3_ar_fusion.finalize_push_fits(num_tokens) + and not use_latent_tail ) current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) @@ -1253,7 +1450,11 @@ def _forward_fused( else: self._forward_routed(hidden_states, router_logits, routed_input, latent) with torch.cuda.stream(self.alt_stream): - self._forward_shared(gate_up, shared_output) + self._forward_shared( + gate_up, + shared_output, + preactivated=shared_is_preactivated, + ) # low-SM pull so the side-stream AR leaves the SMs to the # routed GEMMs it overlaps (K3 dims are fixed; tuned here) k3_ar_fusion.all_reduce_low_sm(shared_output, num_blocks=4, unroll=8) @@ -1273,12 +1474,15 @@ def _forward_fused( *self._get_fused_norm_params(), ) elif self.fuse_ar_norm: - fused_norm = True - k3_ar_fusion.all_reduce_norm( - latent.view(-1, self.moe_hidden_size), - *self._get_fused_norm_params(), - num_tokens=num_tokens, - ) + if use_latent_tail: + k3_ar_fusion.all_reduce(latent) + else: + fused_norm = True + k3_ar_fusion.all_reduce_norm( + latent.view(-1, self.moe_hidden_size), + *self._get_fused_norm_params(), + num_tokens=num_tokens, + ) else: k3_ar_fusion.all_reduce(latent) # the gemm_ag tail wants the normed latent straight out of the @@ -1295,9 +1499,13 @@ def _forward_fused( prefix_sum, ) else: # single collective over the flat [latent | shared] pair - self._forward_shared(gate_up, shared_output) + self._forward_shared( + gate_up, + shared_output, + preactivated=shared_is_preactivated, + ) self._forward_routed(hidden_states, router_logits, routed_input, latent) - if self.fuse_ar_norm and k3_ar_fusion.enabled(): + if self.fuse_ar_norm and k3_ar_fusion.enabled() and not use_latent_tail: fused_norm = True k3_ar_fusion.all_reduce_norm( buf.view(-1, k3_ar_fusion.NORM_DIM), @@ -1311,6 +1519,27 @@ def _forward_fused( latent = buf[:latent_numel].view(num_tokens, self.moe_hidden_size) shared_output = buf[latent_numel:].view(num_tokens, hidden_size) + if use_latent_tail and not fused_norm: + from sglang.kernels.ops.kimi_k3 import latent_tail_aiter_hip + + norm_weight, epsilon = self._get_fused_norm_params() + if latent_tail_aiter_hip.covered( + latent, + shared_output, + norm_weight, + self._latent_tail_weight, + self._latent_tail_scale, + epsilon, + ): + out = latent_tail_aiter_hip.run( + latent, + shared_output, + norm_weight, + self._latent_tail_weight, + self._latent_tail_scale, + epsilon, + ) + return out if prefix_sum is None else out + prefix_sum if not fused_norm: latent = self._latent_norm(latent) out, _ = self.routed_expert_up_proj(latent) @@ -1675,6 +1904,8 @@ def _a_log_weight_loader( self.attn.lower_bound = config.linear_attn_config.get("gate_lower_bound", None) # Set by _prepare_fused_decode() once weights are loaded. self._kda_fused_decode_ready = False + self._kda_group64_weight = None + self._kda_group64_scale = None def forward_qkvbfg(self, hidden_states: torch.Tensor): qkv, _ = self.qkv_proj(hidden_states) @@ -1715,6 +1946,31 @@ def _merge_bfa_weights(self) -> None: self._bfa_f_b_w = self.f_b_proj.weight self._bfa_fa_size, self._bfa_b_size = sizes + def _prepare_group64_projection(self) -> None: + if ( + not _aiter_kda_group64 + or not self.do_fuse_qkvbfg + or not self.use_full_rank_gate + ): + return + from sglang.kernels.ops.kimi_k3 import kda_group64_aiter_hip + + merged = torch.cat( + [ + self.fused_qkvg_proj.weight, + self.b_proj.weight, + self.f_a_proj.weight, + self.f_a_proj.weight.new_zeros((4, self.hidden_size)), + ], + dim=0, + ).contiguous() + if tuple(merged.shape) != (6288, 7168): + return + weight, scale = kda_group64_aiter_hip.pack(merged) + self._kda_group64_weight = weight + self._kda_group64_scale = scale + kda_group64_aiter_hip.warmup(weight, scale) + def _prepare_fused_decode(self) -> None: """Static inputs for the fused KDA decode kernel (kernels/ops/attention/kda_fused_decode): per-segment transposed fp32 conv @@ -1772,6 +2028,31 @@ def _prepare_fused_decode(self) -> None: def forward_qkvbfg_fused(self, hidden_states: torch.Tensor): if self.use_full_rank_gate: + if ( + getattr(self, "_kda_group64_weight", None) is not None + and getattr(self, "_kda_group64_scale", None) is not None + ): + from sglang.kernels.ops.kimi_k3 import ( + kda_group64_aiter_hip, + ) + + if kda_group64_aiter_hip.covered( + hidden_states, + self._kda_group64_weight, + self._kda_group64_scale, + ): + packed = kda_group64_aiter_hip.run( + hidden_states, + self._kda_group64_weight, + self._kda_group64_scale, + ) + mixed_qkv, g_proj_states, beta, f_a, _padding = torch.split( + packed, + [self.split_sizes[0], self.split_sizes[1], 12, 128, 4], + dim=-1, + ) + forget_gate = self.f_b_proj(f_a)[0] + return mixed_qkv, beta, forget_gate, g_proj_states if self._bfa_w is not None: w = self._bfa_w n_fa, n_b = self._bfa_fa_size, self._bfa_b_size @@ -1998,19 +2279,31 @@ def _gated_o_proj_forward(x, *args, **kwargs): # join across graph-segment boundaries. torch.cuda.current_stream().wait_stream(precomputed[1]) if gate_input is not None and not isinstance(x, tuple): - gate = ( - precomputed[0] - if precomputed is not None - else self.g_proj(gate_input)[0] + from sglang.kernels.ops.kimi_k3 import ( + mla_gate_aiter_hip, + mla_output_gate, ) - from sglang.kernels.ops.kimi_k3 import mla_output_gate - if mla_output_gate.covered(x, gate): - # One kernel for x * sigmoid(gate); double rounding - # matches the unfused pair bit-for-bit. - x = mla_output_gate.kimi_k3_mla_output_gate(x, gate) + if ( + precomputed is None + and _aiter_mla_gate + and mla_gate_aiter_hip.covered( + gate_input, self.g_proj.weight, x + ) + ): + x = mla_gate_aiter_hip.run(gate_input, self.g_proj.weight, x) else: - x = x * torch.sigmoid(gate) + gate = ( + precomputed[0] + if precomputed is not None + else self.g_proj(gate_input)[0] + ) + if mla_output_gate.covered(x, gate): + # One kernel for x * sigmoid(gate); double rounding + # matches the unfused pair bit-for-bit. + x = mla_output_gate.kimi_k3_mla_output_gate(x, gate) + else: + x = x * torch.sigmoid(gate) return _orig_o_proj_forward(x, *args, **kwargs) self.o_proj.forward = _gated_o_proj_forward @@ -2023,7 +2316,8 @@ def _precompute_output_gate(self, hidden_states: torch.Tensor) -> None: so its memory cannot be reused while the alt stream still writes.""" self._gate_precomputed = None if ( - self._gate_alt_stream is not None + not _aiter_mla_gate + and self._gate_alt_stream is not None and get_is_capture_mode() # The attention-core break ends the segment between the alt-stream # event record and the o_proj-side wait, so under breakable capture @@ -3073,6 +3367,10 @@ def post_load_weights(self): self_attn.w_vc = w_vc.contiguous().transpose(1, 2) if hasattr(self_attn.kv_b_proj, "weight_scale"): self_attn.w_scale = self_attn.kv_b_proj.weight_scale + if _aiter_mla_gate and isinstance(self_attn, KimiK3MLAAttention): + from sglang.kernels.ops.kimi_k3 import mla_gate_aiter_hip + + mla_gate_aiter_hip.warmup(self_attn.g_proj.weight) # Post-load: precompute the attn-res combined score weights BEFORE # cuda graph capture (a lazy first call inside get_cw would bake the @@ -3100,6 +3398,8 @@ def _warm_cw(proj, norm): continue if isinstance(layer.mlp, KimiK3MoE): layer.mlp._merge_front_weights() + layer.mlp._prepare_preroute_fp8() + layer.mlp._prepare_latent_tail_fp8() # The router consumes the correction bias in fp32; convert the # bf16 checkpoint values once (exact) so the per-call # .to(float32) in topk becomes a no-op instead of one upcast @@ -3109,6 +3409,7 @@ def _warm_cw(proj, norm): bias.data = bias.data.to(torch.float32) if isinstance(layer.self_attn, KimiK3DeltaAttention): layer.self_attn._merge_bfa_weights() + layer.self_attn._prepare_group64_projection() layer.self_attn._prepare_fused_decode() for layer in self.model.layers: diff --git a/test/registered/kernels/ops/kimi_k3/flydsl_ops/__init__.py b/test/registered/kernels/ops/kimi_k3/flydsl_ops/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/test/registered/kernels/ops/kimi_k3/flydsl_ops/bench_kimi_k3_moe_preroute_fp8.py b/test/registered/kernels/ops/kimi_k3/flydsl_ops/bench_kimi_k3_moe_preroute_fp8.py new file mode 100644 index 000000000000..815496887fb8 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/flydsl_ops/bench_kimi_k3_moe_preroute_fp8.py @@ -0,0 +1,121 @@ +"""Graph microbenchmark for Kimi-K3 preroute projections.""" + +import argparse +import statistics + +import torch + +from sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_moe_preroute_fp8 import ( + kimi_k3_moe_tri_projection_cooperative_preactivated_fp8, + kimi_k3_moe_tri_projection_fp8, +) +from sglang.test.ci.ci_register import register_amd_ci +from sglang.utils import is_in_ci + +register_amd_ci(est_time=120, stage="jit-kernel-benchmark", runner_config="amd") + + +def quantize_rows(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + value = weight.float() + scale = value.abs().amax(dim=1).clamp_min(1e-12) / 448.0 + return ( + (value / scale[:, None]).clamp(-448, 448).to(torch.float8_e4m3fn).contiguous(), + scale.contiguous(), + ) + + +def bench(fn, warmup: int, iters: int, trials: int) -> tuple[float, float, float]: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fn() + for _ in range(warmup): + graph.replay() + torch.cuda.synchronize() + samples = [] + for _ in range(trials): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000 / iters) + ordered = sorted(samples) + return ( + statistics.median(samples), + ordered[max(0, trials // 10 - 1)], + ordered[min(trials - 1, 9 * trials // 10)], + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", type=int, choices=(1, 2, 4), required=True) + parser.add_argument("--warmup", type=int, default=50) + parser.add_argument("--iters", type=int, default=500) + parser.add_argument("--trials", type=int, default=11) + args = parser.parse_args() + torch.manual_seed(20260817) + device = torch.device("cuda") + hidden = torch.randn((args.tokens, 7168), dtype=torch.bfloat16, device=device) + routed_bf16 = torch.randn((3584, 7168), dtype=torch.bfloat16, device=device) + shared_bf16 = torch.randn((1536, 7168), dtype=torch.bfloat16, device=device) + router = torch.randn((896, 7168), dtype=torch.bfloat16, device=device) + routed, routed_scale = quantize_rows(routed_bf16) + shared, shared_scale = quantize_rows(shared_bf16) + merged = torch.cat((shared_bf16, router, routed_bf16), dim=0).contiguous() + if args.tokens in (2, 4): + shared = ( + shared.view(2, 768, 7168).permute(1, 0, 2).contiguous().view(1536, 7168) + ) + shared_scale = shared_scale.view(2, 768).t().contiguous().view(1536) + + def rowdot_fn(): + return kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + hidden, + routed, + routed_scale, + shared, + shared_scale, + router, + situ_beta=4.0, + situ_linear_beta=25.0, + fast_situ=True, + ) + + else: + + def rowdot_fn(): + return kimi_k3_moe_tri_projection_fp8( + hidden, routed, routed_scale, shared, shared_scale, router + ) + + rowdot = bench( + rowdot_fn, + args.warmup, + args.iters, + args.trials, + ) + bf16 = bench( + lambda: torch.mm(hidden, merged.t()), + args.warmup, + args.iters, + args.trials, + ) + print( + f"M={args.tokens} rowdot_p50_us={rowdot[0]:.3f} " + f"rowdot_p10_us={rowdot[1]:.3f} rowdot_p90_us={rowdot[2]:.3f} " + f"bf16_p50_us={bf16[0]:.3f} " + f"rowdot_speedup={bf16[0] / rowdot[0]:.3f}x" + ) + + +if __name__ == "__main__": + if is_in_ci(): + print("Skipping bench_kimi_k3_moe_preroute_fp8.py in CI") + raise SystemExit(0) + main() diff --git a/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_kda_input_group64.py b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_kda_input_group64.py new file mode 100644 index 000000000000..5d1091951ea5 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_kda_input_group64.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import importlib + +import pytest +import torch +from aiter.jit.utils.chip_info import get_gfx + +from sglang.kernels.ops.kimi_k3.flydsl.kernels.kimi_k3_kda_input_group64_gfx950 import ( + build_kimi_k3_kda_input_group64_module, +) +from sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_kda_input_group64 import ( + kimi_k3_kda_input_group64, + quantize_kimi_k3_kda_input_group64, + supports_kimi_k3_kda_input_group64, +) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x") + +group64_module = importlib.import_module( + "sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_kda_input_group64" +) + + +def test_support_predicate_fails_closed_off_gpu() -> None: + tensor = torch.empty(1) + assert not supports_kimi_k3_kda_input_group64(tensor, tensor, tensor) + + +def test_quantizer_rejects_non_cuda_input() -> None: + with pytest.raises(ValueError, match="contiguous CUDA BF16"): + quantize_kimi_k3_kda_input_group64(torch.empty(1, dtype=torch.bfloat16)) + + +def test_wrapper_uses_validated_gfx950_schedule_by_default(monkeypatch) -> None: + hidden = torch.empty((1, 7168), dtype=torch.bfloat16) + weight = torch.empty((6284, 7168), dtype=torch.float8_e4m3fn) + scale = torch.empty((6284, 112), dtype=torch.float32) + output = torch.empty((1, 6288), dtype=torch.bfloat16) + schedules = [] + launches = [] + + def fake_builder(**kwargs): + schedules.append(kwargs) + + def launch(*args, **kwargs): + launches.append((args, kwargs)) + + return launch + + monkeypatch.setattr( + group64_module, + "supports_kimi_k3_kda_input_group64", + lambda *args: True, + ) + monkeypatch.setattr( + group64_module, + "build_kimi_k3_kda_input_group64_module", + fake_builder, + ) + monkeypatch.setattr(group64_module, "ptr_arg", lambda tensor: tensor) + monkeypatch.setattr(torch.cuda, "current_stream", lambda device: None) + group64_module._launcher.cache_clear() + + assert kimi_k3_kda_input_group64(hidden, weight, scale, output) is output + assert schedules == [ + { + "num_tokens": 1, + "rows_per_wave": 2, + "cu_count": 256, + "waves_per_eu": 0, + "weight_cache_modifier": 2, + "hidden_to_lds": True, + } + ] + assert len(launches) == 1 + group64_module._launcher.cache_clear() + + +@pytest.mark.parametrize("rows_per_wave", [0, 5]) +def test_builder_rejects_unsupported_rows_per_wave(rows_per_wave: int) -> None: + with pytest.raises(ValueError, match="rows_per_wave"): + build_kimi_k3_kda_input_group64_module(rows_per_wave=rows_per_wave) + + +@pytest.mark.parametrize("cu_count", [0, 257]) +def test_builder_rejects_unsupported_cu_count(cu_count: int) -> None: + with pytest.raises(ValueError, match="cu_count"): + build_kimi_k3_kda_input_group64_module(cu_count=cu_count) + + +def test_builder_rejects_negative_waves_per_eu() -> None: + with pytest.raises(ValueError, match="waves_per_eu"): + build_kimi_k3_kda_input_group64_module(waves_per_eu=-1) + + +@pytest.mark.parametrize("cache_modifier", [-1, 4]) +def test_builder_rejects_unsupported_cache_modifier( + cache_modifier: int, +) -> None: + with pytest.raises(ValueError, match="weight_cache_modifier"): + build_kimi_k3_kda_input_group64_module(weight_cache_modifier=cache_modifier) + + +@pytest.mark.parametrize("num_tokens", [1, 2]) +@torch.inference_mode() +def test_group64_projection_matches_dequantized_reference_on_gfx950( + num_tokens: int, +) -> None: + if not torch.cuda.is_available() or get_gfx() != "gfx950": + pytest.skip("Kimi-K3 group64 projection is gfx950-only") + + hidden_size = 7168 + stored_rows = 6284 + output_rows = 6288 + group_size = 64 + groups_per_row = hidden_size // group_size + + columns = torch.arange(hidden_size, device="cuda", dtype=torch.float32) + rows = torch.arange(stored_rows, device="cuda", dtype=torch.float32) + hidden = torch.stack( + [ + torch.sin(columns * (0.03125 + token * 0.001)).to(torch.bfloat16) + for token in range(num_tokens) + ] + ) + stored_weight = ( + ((columns.to(torch.int32) % 17).float() - 8.0).unsqueeze(0) + * ((rows.to(torch.int32) % 13).float() + 1.0).unsqueeze(1) + * 0.002 + ).to(torch.bfloat16) + weight = torch.zeros( + output_rows, + hidden_size, + device="cuda", + dtype=torch.bfloat16, + ) + weight[:stored_rows].copy_(stored_weight) + + packed, scale = quantize_kimi_k3_kda_input_group64(weight) + actual = kimi_k3_kda_input_group64(hidden, packed, scale) + + hidden_groups = hidden.float().reshape(num_tokens, groups_per_row, group_size) + dequantized = ( + packed.to(torch.bfloat16) + .float() + .reshape(stored_rows, groups_per_row, group_size) + * scale[..., None] + ) + expected = torch.einsum("tgc,rgc->tr", hidden_groups, dequantized).to( + torch.bfloat16 + ) + actual_fp32 = actual[:, :stored_rows].float() + expected_fp32 = expected.float() + delta = actual_fp32 - expected_fp32 + relative_rmse = delta.square().mean().sqrt() / expected_fp32.square().mean().sqrt() + cosine = torch.nn.functional.cosine_similarity( + actual_fp32, + expected_fp32, + dim=1, + ) + + assert torch.isfinite(actual).all() + assert relative_rmse <= 5e-4 + assert torch.all(cosine >= 0.99999) + assert torch.count_nonzero(actual[:, stored_rows:]).item() == 0 + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_mla_gate_epilogue.py b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_mla_gate_epilogue.py new file mode 100644 index 000000000000..3dd6b6e0b4f7 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_mla_gate_epilogue.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +from sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_mla_gate import ( + kimi_k3_mla_gate, + supports_kimi_k3_mla_gate, +) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x") + + +def _gfx950_flydsl_available() -> bool: + if not torch.cuda.is_available() or not is_flydsl_available(): + return False + try: + return get_gfx_runtime() == "gfx950" + except (AssertionError, KeyError, RuntimeError): + return False + + +pytestmark = pytest.mark.skipif( + not _gfx950_flydsl_available(), + reason="Kimi-K3 MLA gate specialization requires gfx950", +) + + +def _inputs(batch: int, seed: int = 17): + torch.manual_seed(seed) + hidden = torch.randn((batch, 7168), device="cuda", dtype=torch.bfloat16) * 0.15 + weight = torch.randn((1536, 7168), device="cuda", dtype=torch.bfloat16) * 0.02 + attention = torch.randn((batch, 1536), device="cuda", dtype=torch.bfloat16) + return hidden.contiguous(), weight.contiguous(), attention.contiguous() + + +def _relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float: + error = actual.float() - expected.float() + return ( + torch.sqrt(torch.mean(error.square())) + / torch.sqrt(torch.mean(expected.float().square())).clamp_min(1e-30) + ).item() + + +def _reference( + hidden: torch.Tensor, + weight: torch.Tensor, + attention: torch.Tensor, +) -> torch.Tensor: + projected = (hidden.float() @ weight.float().T).to(torch.bfloat16) + gate = torch.sigmoid(projected.float()).to(torch.bfloat16) + return (gate.float() * attention.float()).to(torch.bfloat16) + + +def test_kimi_k3_mla_gate_primary_dispatch_and_accuracy(): + hidden, weight, attention = _inputs(1) + assert supports_kimi_k3_mla_gate(hidden, weight, attention) + output = kimi_k3_mla_gate(hidden, weight, attention) + reference = _reference(hidden, weight, attention) + assert output.dtype == torch.bfloat16 + assert output.is_contiguous() + assert _relative_rmse(output, reference) <= 0.01 + assert ( + torch.nn.functional.cosine_similarity( + output.float().flatten(), + reference.float().flatten(), + dim=0, + ).item() + >= 0.999 + ) + + +def test_kimi_k3_mla_gate_reuses_valid_output(): + hidden, weight, attention = _inputs(1) + output = torch.empty_like(attention) + + actual = kimi_k3_mla_gate(hidden, weight, attention, out=output) + torch.cuda.synchronize() + + assert actual is output + + +def test_kimi_k3_mla_gate_supports_attention_output_alias(): + hidden, weight, attention = _inputs(1) + attention_input = attention.clone() + reference = _reference(hidden, weight, attention_input) + + actual = kimi_k3_mla_gate(hidden, weight, attention, out=attention) + torch.cuda.synchronize() + + assert actual is attention + assert _relative_rmse(actual, reference) <= 0.01 + + +def test_kimi_k3_mla_gate_graph_replay_uses_changed_inputs(): + hidden, weight, attention = _inputs(1) + output = torch.empty_like(attention) + kimi_k3_mla_gate(hidden, weight, attention, out=output) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + kimi_k3_mla_gate(hidden, weight, attention, out=output) + + graph.replay() + torch.cuda.synchronize() + first = output.clone() + hidden.copy_(torch.randn_like(hidden)) + attention.copy_(torch.randn_like(attention)) + graph.replay() + torch.cuda.synchronize() + second = output.clone() + + assert not torch.equal(first, second) + reference = _reference(hidden, weight, attention) + assert _relative_rmse(second, reference) <= 0.01 + + +def test_kimi_k3_mla_gate_support_is_narrow(): + hidden, weight, attention = _inputs(1) + noncontiguous_attention = torch.empty( + (1, 3072), + device="cuda", + dtype=torch.bfloat16, + )[:, ::2] + + assert supports_kimi_k3_mla_gate(hidden, weight, attention) + assert not supports_kimi_k3_mla_gate(hidden.expand(2, -1), weight, attention) + assert not supports_kimi_k3_mla_gate(hidden, weight, noncontiguous_attention) + + +def test_kimi_k3_mla_gate_rejects_unsupported_shape(): + hidden, weight, attention = _inputs(2) + + with pytest.raises(NotImplementedError, match="requires contiguous gfx950 BF16"): + kimi_k3_mla_gate(hidden, weight, attention) + + +@pytest.mark.parametrize( + "shape,dtype", + [ + ((1536,), torch.bfloat16), + ((1, 1536), torch.float32), + ], +) +def test_kimi_k3_mla_gate_rejects_invalid_output(shape, dtype): + hidden, weight, attention = _inputs(1) + output = torch.empty(shape, device="cuda", dtype=dtype) + + with pytest.raises(ValueError, match="out must be contiguous BF16"): + kimi_k3_mla_gate(hidden, weight, attention, out=output) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_moe_preroute_fp8.py b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_moe_preroute_fp8.py new file mode 100644 index 000000000000..96f159ff91a7 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_kimi_k3_moe_preroute_fp8.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch +import torch.nn.functional as F +from aiter.jit.utils.chip_info import get_gfx_runtime +from aiter.ops.flydsl.utils import is_flydsl_available + +from sglang.kernels.ops.kimi_k3.flydsl.kimi_k3_moe_preroute_fp8 import ( + is_kimi_k3_moe_preroute_fp8_available, + kimi_k3_moe_dual_projection_fp8, + kimi_k3_moe_tri_projection_cooperative_preactivated_fp8, + kimi_k3_moe_tri_projection_fp8, + kimi_k3_shared_down_fp8, + supports_kimi_k3_moe_dual_projection_fp8, + supports_kimi_k3_moe_tri_projection_cooperative_preactivated_fp8, + supports_kimi_k3_moe_tri_projection_fp8, + supports_kimi_k3_shared_down_fp8, + supports_kimi_k3_shared_down_fp8_weight, +) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=120, suite="stage-b-test-1-gpu-small-amd-mi35x") + +_FP8_MAX = 448.0 + + +def _quantize_rows( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + weight_f32 = weight.float() + amax = weight_f32.abs().amax(dim=1) + scale = torch.where( + amax > 0, + amax / _FP8_MAX, + torch.ones_like(amax), + ) + quantized = ( + (weight_f32 / scale[:, None]) + .clamp(min=-_FP8_MAX, max=_FP8_MAX) + .to(torch.float8_e4m3fn) + .contiguous() + ) + return quantized, scale.contiguous() + + +def _relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float: + error = (actual.float() - expected.float()).square().mean().sqrt() + reference = expected.float().square().mean().sqrt().clamp_min(1e-12) + return (error / reference).item() + + +def test_support_predicates_fail_closed_on_cpu(): + hidden = torch.empty((1, 7168), dtype=torch.bfloat16) + fp8 = torch.empty((1, 1), dtype=torch.float8_e4m3fn) + scale = torch.empty((1,), dtype=torch.float32) + + assert not supports_kimi_k3_moe_dual_projection_fp8( + hidden, + fp8, + scale, + fp8, + scale, + ) + assert not supports_kimi_k3_moe_tri_projection_fp8( + hidden, + fp8, + scale, + fp8, + scale, + hidden, + ) + assert not supports_kimi_k3_shared_down_fp8(hidden, fp8, scale) + assert not supports_kimi_k3_shared_down_fp8_weight(fp8, scale) + assert not supports_kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + torch.empty((4, 7168), dtype=torch.bfloat16), + fp8, + scale, + fp8, + scale, + hidden, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a GPU runtime") +def test_backend_availability_matches_flydsl_and_architecture(): + assert is_kimi_k3_moe_preroute_fp8_available() == ( + is_flydsl_available() and get_gfx_runtime() == "gfx950" + ) + + +@pytest.mark.parametrize( + ("situ_beta", "situ_linear_beta"), + [(0.0, 25.0), (4.0, -1.0), (float("nan"), 25.0)], +) +def test_shared_down_rejects_invalid_situ_parameters( + situ_beta: float, + situ_linear_beta: float, +): + tensor = torch.empty((1,), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="finite and positive"): + kimi_k3_shared_down_fp8( + tensor, + tensor, + tensor.float(), + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() + or not is_flydsl_available() + or get_gfx_runtime() != "gfx950", + reason="requires FlyDSL on gfx950", +) +def test_kimi_k3_preroute_fp8_matches_dequantized_reference(): + torch.manual_seed(20260729) + device = torch.device("cuda") + hidden = torch.randn((1, 7168), device=device, dtype=torch.bfloat16) + routed_bf16 = torch.randn( + (3584, 7168), + device=device, + dtype=torch.bfloat16, + ) + shared_up_bf16 = torch.randn( + (1536, 7168), + device=device, + dtype=torch.bfloat16, + ) + shared_down_bf16 = torch.randn( + (7168, 768), + device=device, + dtype=torch.bfloat16, + ) + router_bf16 = torch.randn( + (896, 7168), + device=device, + dtype=torch.bfloat16, + ) + routed_weight, routed_scale = _quantize_rows(routed_bf16) + shared_up_weight, shared_up_scale = _quantize_rows(shared_up_bf16) + shared_down_weight, shared_down_scale = _quantize_rows(shared_down_bf16) + + routed, gate_up = kimi_k3_moe_dual_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_up_weight, + shared_up_scale, + ) + tri_routed, tri_gate_up, router_logits = kimi_k3_moe_tri_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_up_weight, + shared_up_scale, + router_bf16, + ) + shared = kimi_k3_shared_down_fp8( + gate_up, + shared_down_weight, + shared_down_scale, + situ_beta=4.0, + situ_linear_beta=25.0, + ) + + routed_dequant = routed_weight.float() * routed_scale[:, None] + shared_up_dequant = shared_up_weight.float() * shared_up_scale[:, None] + shared_down_dequant = shared_down_weight.float() * shared_down_scale[:, None] + routed_ref = hidden.float() @ routed_dequant.t() + gate_up_ref = hidden.float() @ shared_up_dequant.t() + gate, up = gate_up_ref.to(torch.bfloat16).float().chunk(2, dim=-1) + activated = ( + ( + 4.0 + * torch.tanh(gate / 4.0) + * torch.sigmoid(gate) + * 25.0 + * torch.tanh(up / 25.0) + ) + .to(torch.bfloat16) + .float() + ) + shared_ref = activated @ shared_down_dequant.t() + + assert _relative_rmse(routed, routed_ref) < 0.035 + assert _relative_rmse(gate_up, gate_up_ref) < 0.035 + # fdot2 may round in a different reduction order than the dual kernel. + assert _relative_rmse(tri_routed, routed_ref) < 0.035 + assert _relative_rmse(tri_gate_up, gate_up_ref) < 0.035 + router_ref = F.linear(hidden, router_bf16).float() + assert _relative_rmse(router_logits, router_ref) < 0.01 + router_topk = router_logits.topk(17, dim=-1) + reference_topk = router_ref.topk(17, dim=-1) + assert reference_topk.values[0, 15] > reference_topk.values[0, 16] + # torch.topk does not define the order of equal values within the result. + torch.testing.assert_close( + router_topk.indices[:, :16].sort(dim=-1).values, + reference_topk.indices[:, :16].sort(dim=-1).values, + atol=0, + rtol=0, + ) + assert _relative_rmse(shared, shared_ref) < 0.06 + assert F.cosine_similarity(shared.float(), shared_ref.float()).item() > 0.998 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_routed, captured_gate_up = kimi_k3_moe_dual_projection_fp8( + hidden, + routed_weight, + routed_scale, + shared_up_weight, + shared_up_scale, + ) + captured_shared = kimi_k3_shared_down_fp8( + captured_gate_up, + shared_down_weight, + shared_down_scale, + situ_beta=4.0, + situ_linear_beta=25.0, + ) + graph.replay() + expected_routed = captured_routed.clone() + expected_shared = captured_shared.clone() + graph.replay() + torch.testing.assert_close(captured_routed, expected_routed, atol=0, rtol=0) + torch.testing.assert_close(captured_shared, expected_shared, atol=0, rtol=0) + + +@pytest.mark.parametrize("num_tokens", [2, 4]) +@pytest.mark.skipif( + not torch.cuda.is_available() + or not is_flydsl_available() + or get_gfx_runtime() != "gfx950", + reason="requires FlyDSL on gfx950", +) +def test_kimi_k3_cooperative_preactivated_projection(num_tokens: int): + torch.manual_seed(20260818 + num_tokens) + hidden = torch.randn((num_tokens, 7168), device="cuda", dtype=torch.bfloat16) + routed_bf16 = torch.randn((3584, 7168), device="cuda", dtype=torch.bfloat16) + shared_bf16 = torch.randn((1536, 7168), device="cuda", dtype=torch.bfloat16) + router_bf16 = torch.randn((896, 7168), device="cuda", dtype=torch.bfloat16) + routed_weight, routed_scale = _quantize_rows(routed_bf16) + shared_weight, shared_scale = _quantize_rows(shared_bf16) + shared_interleaved = ( + shared_weight.view(2, 768, 7168).permute(1, 0, 2).contiguous().view(1536, 7168) + ) + shared_scale_interleaved = shared_scale.view(2, 768).t().contiguous().view(1536) + + routed, activated, router = kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + hidden, + routed_weight, + routed_scale, + shared_interleaved, + shared_scale_interleaved, + router_bf16, + situ_beta=4.0, + situ_linear_beta=25.0, + fast_situ=True, + ) + routed_ref = hidden.float() @ (routed_weight.float() * routed_scale[:, None]).t() + gate_up_ref = ( + hidden.float() @ (shared_weight.float() * shared_scale[:, None]).t() + ).to(torch.bfloat16) + gate, up = gate_up_ref.float().chunk(2, dim=-1) + activated_ref = ( + 4.0 + * torch.tanh(gate / 4.0) + * torch.sigmoid(gate) + * 25.0 + * torch.tanh(up / 25.0) + ).to(torch.bfloat16) + router_ref = F.linear(hidden, router_bf16).float() + + assert _relative_rmse(routed, routed_ref) < 0.035 + assert _relative_rmse(activated, activated_ref) < 0.035 + assert _relative_rmse(router, router_ref) < 0.01 + torch.testing.assert_close( + router.topk(16, dim=-1).indices.sort(dim=-1).values, + router_ref.topk(16, dim=-1).indices.sort(dim=-1).values, + atol=0, + rtol=0, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = kimi_k3_moe_tri_projection_cooperative_preactivated_fp8( + hidden, + routed_weight, + routed_scale, + shared_interleaved, + shared_scale_interleaved, + router_bf16, + situ_beta=4.0, + situ_linear_beta=25.0, + fast_situ=True, + ) + graph.replay() + expected = tuple(value.clone() for value in captured) + graph.replay() + for actual, reference in zip(captured, expected): + torch.testing.assert_close(actual, reference, atol=0, rtol=0) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_latent_moe_tail_fp8.py b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_latent_moe_tail_fp8.py new file mode 100644 index 000000000000..8c9a487082e9 --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/flydsl_ops/test_latent_moe_tail_fp8.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import pytest +import torch +from aiter.jit.utils.chip_info import get_gfx_runtime + +from sglang.kernels.ops.kimi_k3.flydsl.kernels.latent_moe_tail_fp8_gfx950 import ( + build_b1_latent_moe_tail_fp8_persistent_module, +) +from sglang.kernels.ops.kimi_k3.flydsl.latent_moe_tail_fp8 import ( + latent_moe_tail_fp8, + quantize_latent_moe_tail_weight, + supports_latent_moe_tail_fp8, +) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd-mi35x") + +LATENT_DIM = 3584 +HIDDEN_DIM = 7168 +EPSILON = 1.0e-6 + + +def _relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float: + error = (actual.float() - expected.float()).square().mean().sqrt() + reference = expected.float().square().mean().sqrt().clamp_min(1.0e-12) + return (error / reference).item() + + +def test_support_predicate_fails_closed_on_cpu(): + bf16 = torch.empty(1, dtype=torch.bfloat16) + fp8 = torch.empty(1, dtype=torch.float8_e4m3fn) + fp32 = torch.empty(1, dtype=torch.float32) + + assert not supports_latent_moe_tail_fp8( + bf16, + bf16, + bf16, + fp8, + fp32, + EPSILON, + ) + + +def test_quantizer_rejects_non_cuda_input(): + with pytest.raises(ValueError, match="contiguous CUDA BF16"): + quantize_latent_moe_tail_weight(torch.empty(1, dtype=torch.bfloat16)) + + +@pytest.mark.parametrize( + ("keyword", "value", "message"), + [ + ("rows_per_wave", 0, "rows_per_wave"), + ("cu_count", 257, "cu_count"), + ("waves_per_eu", -1, "waves_per_eu"), + ("weight_cache_modifier", 4, "weight_cache_modifier"), + ], +) +def test_builder_rejects_invalid_schedule(keyword: str, value: int, message: str): + with pytest.raises(ValueError, match=message): + build_b1_latent_moe_tail_fp8_persistent_module(**{keyword: value}) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or get_gfx_runtime() != "gfx950", + reason="Kimi-K3 FP8 latent-tail specialization requires gfx950", +) +@torch.inference_mode() +def test_latent_moe_tail_fp8_matches_dequantized_oracle_and_replays(): + generator = torch.Generator(device="cpu").manual_seed(20260730) + routed = torch.randn((1, LATENT_DIM), generator=generator).bfloat16().cuda() + shared = torch.randn((1, HIDDEN_DIM), generator=generator).bfloat16().cuda() + rms_weight = torch.randn(LATENT_DIM, generator=generator).bfloat16().cuda() + up_weight = ( + torch.randn((HIDDEN_DIM, LATENT_DIM), generator=generator) + .mul_(LATENT_DIM**-0.5) + .bfloat16() + .cuda() + ) + packed, scale = quantize_latent_moe_tail_weight(up_weight) + + inverse_rms = torch.rsqrt( + routed.float().square().mean(dim=-1, keepdim=True) + EPSILON + ) + normalized = (routed.float() * inverse_rms * rms_weight.float()).bfloat16() + dequantized = packed.float() * scale[:, None] + dequantized_oracle = ( + torch.mm(normalized.float(), dequantized.t()).bfloat16().float() + + shared.float() + ).bfloat16() + bf16_oracle = ( + torch.mm(normalized.float(), up_weight.float().t()).bfloat16().float() + + shared.float() + ).bfloat16() + + out = torch.empty_like(shared) + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + latent_moe_tail_fp8( + routed, + shared, + rms_weight, + packed, + scale, + EPSILON, + out=out, + ) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = latent_moe_tail_fp8( + routed, + shared, + rms_weight, + packed, + scale, + EPSILON, + out=out, + ) + graph.replay() + torch.cuda.synchronize() + + assert actual is out + torch.testing.assert_close(actual, dequantized_oracle, rtol=0.01, atol=0.015625) + assert _relative_rmse(actual, bf16_oracle) < 0.03 + assert ( + torch.nn.functional.cosine_similarity( + actual.float(), + bf16_oracle.float(), + ).item() + > 0.999 + ) + + previous = actual.clone() + routed.add_(0.5) + graph.replay() + torch.cuda.synchronize() + assert not torch.equal(previous, actual) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/kernels/ops/kimi_k3/test_aiter_feature_gates.py b/test/registered/kernels/ops/kimi_k3/test_aiter_feature_gates.py new file mode 100644 index 000000000000..6457f8ecfeea --- /dev/null +++ b/test/registered/kernels/ops/kimi_k3/test_aiter_feature_gates.py @@ -0,0 +1,57 @@ +import importlib + +import pytest + +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=5, suite="stage-b-test-1-gpu-small-amd-mi35x") + +pytestmark = pytest.mark.skipif(not is_hip(), reason="Kimi-K3 AITER gates require HIP") + + +@pytest.mark.parametrize( + ("module_name", "env_name"), + [ + ( + "sglang.kernels.ops.kimi_k3.mla_gate_aiter_hip", + "SGLANG_K3_AITER_MLA_GATE", + ), + ( + "sglang.kernels.ops.kimi_k3.kda_group64_aiter_hip", + "SGLANG_K3_AITER_KDA_GROUP64", + ), + ( + "sglang.kernels.ops.kimi_k3.moe_preroute_aiter_hip", + "SGLANG_K3_AITER_MOE_PREROUTE_FP8", + ), + ( + "sglang.kernels.ops.kimi_k3.latent_tail_aiter_hip", + "SGLANG_K3_AITER_LATENT_TAIL_FP8", + ), + ], +) +def test_aiter_feature_is_opt_in(monkeypatch, module_name, env_name): + module = importlib.import_module(module_name) + monkeypatch.delenv(env_name, raising=False) + assert not module.enabled() + monkeypatch.setenv(env_name, "1") + assert module.enabled() + + +def test_moe_preroute_cooperative_feature_is_opt_in(monkeypatch): + module = importlib.import_module( + "sglang.kernels.ops.kimi_k3.moe_preroute_aiter_hip" + ) + monkeypatch.setenv("SGLANG_K3_AITER_MOE_PREROUTE_FP8", "1") + monkeypatch.setenv("SGLANG_K3_FLYDSL_SOURCE", "sglang") + monkeypatch.delenv("SGLANG_K3_PREROUTE_PREACTIVATED_SHARED", raising=False) + assert not module.cooperative_preactivated_enabled() + monkeypatch.setenv("SGLANG_K3_PREROUTE_PREACTIVATED_SHARED", "1") + assert module.cooperative_preactivated_enabled() + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))