diff --git a/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py b/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py new file mode 100644 index 000000000000..9e4f4157a8a4 --- /dev/null +++ b/benchmarks/kernels/benchmark_flydsl_moe_w4a16.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2025 FlyDSL Project Contributors + +import json +import os + +import torch +from aiter.test_common import run_perftest + +from vllm.model_executor.layers.fused_moe import fused_experts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import fused_flydsl_moe +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + compressed_tensors_moe_w4a16_flydsl, +) +from vllm.platforms import current_platform + +RoutingBuffers = tuple[ + torch.Tensor, # sorted_token_ids + torch.Tensor, # sorted_weights + torch.Tensor, # sorted_expert_ids + torch.Tensor, # num_valid_ids (shape [1], i32) + int, # sorted_size + int, # blocks +] + +MODEL_PARAMS_TO_TUNE = [ + # (num_experts, inter_dim, hidden_size, topk) + (384, 256, 7168, 8), # Kimi K2.5 TP=8 + (384, 512, 7168, 8), # Kimi K2.5 TP=4 +] + +NUM_TOKENS_TO_TUNE = [ + 1, + 2, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, +] + +TILE_M_SEARCH_SPACE = [16, 32, 64, 128, 256] +TILE_N_SEARCH_SPACE = [16, 32, 64, 128, 256] +TILE_K_SEARCH_SPACE = [16, 32, 64, 128, 256, 512] +TILE_N2_SEARCH_SPACE = [16, 32, 64, 128, 256] +TILE_K2_SEARCH_SPACE = [16, 32, 64, 128, 256, 512] + +TILE_CONFIGS = [] +for tile_m in TILE_M_SEARCH_SPACE: + for tile_n in TILE_N_SEARCH_SPACE: + for tile_k in TILE_K_SEARCH_SPACE: + for tile_n2 in TILE_N2_SEARCH_SPACE: + for tile_k2 in TILE_K2_SEARCH_SPACE: + TILE_CONFIGS.append( + { + "tile_m": tile_m, + "tile_n": tile_n, + "tile_k": tile_k, + "tile_n2": tile_n2, + "tile_k2": tile_k2, + } + ) + + +def tune_flydsl_moe_w4a16( + device: str = "cuda", num_iters: int = 100, num_warmup: int = 10 +): + packed_factor = 8 + w13_num_shards = 2 + params_dtype = torch.bfloat16 + group_size = 32 + scale_factor = 0.01 + + for model_params in MODEL_PARAMS_TO_TUNE: + num_experts = model_params[0] + inter_dim = model_params[1] + hidden_size = model_params[2] + topk = model_params[3] + print( + f"\nTuning: num_experts={num_experts}, inter_dim={inter_dim}, " + f"hidden_size={hidden_size}, topk={topk}...\n" + ) + + w2_scales_size = inter_dim + num_groups_w2 = w2_scales_size // group_size + num_groups_w13 = hidden_size // group_size + + w13_weight = torch.randint( + 0, + 255, + (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim), + dtype=torch.int32, + device=device, + ) + + w2_weight = torch.randint( + 0, + 255, + (num_experts, inter_dim // packed_factor, hidden_size), + dtype=torch.int32, + device=device, + ) + w13_scale = scale_factor * torch.randn( + num_experts, + num_groups_w13, + w13_num_shards * inter_dim, + dtype=params_dtype, + device=device, + ) + w2_scale = scale_factor * torch.randn( + num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device + ) + + w13 = w13_weight + w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13) + w13 = w13.view(-1).contiguous() + + w2 = w2_weight + w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2) + w2 = w2.view(-1).contiguous() + + w13_scale_flydsl = w13_scale + w2_scale_flydsl = w2_scale + + if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1: + E, G, N = w13_scale.shape + w13_scale_flydsl = ( + w13_scale_flydsl.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1: + w13_scale_flydsl = w13_scale_flydsl.squeeze(1) + + if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1: + E, G, N = w2_scale.shape + w2_scale_flydsl = ( + w2_scale_flydsl.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1: + w2_scale_flydsl = w2_scale_flydsl.squeeze(1) + + w13_scale_flydsl = w13_scale_flydsl.contiguous() + w2_scale_flydsl = w2_scale_flydsl.contiguous() + + w13.is_shuffled = True + w2.is_shuffled = True + + w13_weight_scale = w13_scale.transpose(1, 2).contiguous() + w2_weight_scale = w2_scale.transpose(1, 2).contiguous() + w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8) + w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8) + + moe_quant_config = int4_w4a16_moe_quant_config( + w1_scale=w13_weight_scale, + w2_scale=w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) + + tuned_config = {} + + for num_tokens in NUM_TOKENS_TO_TUNE: + score = torch.rand( + (num_tokens, num_experts), device=device, dtype=torch.float32 + ) + topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) + topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) + x = torch.randn( + (num_tokens, hidden_size), dtype=torch.bfloat16, device=device + ) + us_best = float("inf") + for tile_config in TILE_CONFIGS: + try: + tile_m = tile_config["tile_m"] + tile_n = tile_config["tile_n"] + tile_k = tile_config["tile_k"] + tile_n2 = tile_config["tile_n2"] + tile_k2 = tile_config["tile_k2"] + + model_dim = x.shape[1] + assert model_dim % 64 == 0 + assert model_dim % tile_k == 0 + assert inter_dim % tile_n == 0 + assert model_dim % tile_n2 == 0 + assert inter_dim % tile_k2 == 0 + assert ((tile_m * tile_k2) % 256) == 0 + bytes_per_thread_x = (tile_m * tile_k2) // 256 + assert (bytes_per_thread_x % 4) == 0 + + out, _us = run_perftest( + fused_flydsl_moe, + x, + w13, + w2, + num_experts, + inter_dim, + topk_weights, + topk_ids, + num_iters=num_iters, + num_warmup=num_warmup, + w1_scale=w13_scale_flydsl, + w2_scale=w2_scale_flydsl, + topk=topk_weights.shape[-1], + group_size=group_size, + doweight_stage1=False, + scale_is_bf16=True, + config=tile_config, + ) + torch.accelerator.synchronize() + except Exception: + torch.accelerator.synchronize() + continue + else: + us = _us.item() + if us < us_best: + out_ref = fused_experts( + x, + w13_weight_packed, + w2_weight_packed, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + global_num_experts=num_experts, + expert_map=None, + quant_config=moe_quant_config, + ) + try: + assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1) + except Exception: + continue + else: + print( + f"For [num_tokens={num_tokens}, num_experts={num_experts}, " # noqa: E501 + f"inter_dim={inter_dim}] found new best " # noqa: E501 + f"config={tile_config}, us={us:0.3f}" + ) + us_best = us + tuned_config[str(num_tokens)] = tile_config + device_name = current_platform.get_device_name().replace(" ", "_") + tuned_config_file_name = ( + f"E={num_experts},N={inter_dim},device_name={device_name}," + f"dtype=int4_w4a16,backend=flydsl.json" + ) + tuner_dir_path = os.path.dirname(os.path.realpath(__file__)) + store_path = os.path.join(tuner_dir_path, tuned_config_file_name) + with open(store_path, "w") as f: + json.dump(tuned_config, f, indent=4) + print( + f"\nTuned config for num_tokens={num_tokens} was stored at {store_path}\n" # noqa: E501 + ) + + +if __name__ == "__main__": + tune_flydsl_moe_w4a16(device="cuda") diff --git a/tests/kernels/moe/test_flydsl_moe.py b/tests/kernels/moe/test_flydsl_moe.py new file mode 100644 index 000000000000..7c51c3691311 --- /dev/null +++ b/tests/kernels/moe/test_flydsl_moe.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright (c) 2025 FlyDSL Project Contributors + + +import importlib.util + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe import fused_experts +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + int4_w4a16_moe_quant_config, +) +from vllm.platforms import current_platform +from vllm.platforms.rocm import on_gfx950 + +if not (current_platform.is_rocm() and on_gfx950()): + pytest.skip("This test can only run on ROCm and gfx950.", allow_module_level=True) + +aiter_available = importlib.util.find_spec("aiter") is not None + +if not aiter_available: + pytest.skip("These tests require AITER to run.", allow_module_level=True) + +from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import ( # noqa: E402 + fused_flydsl_moe, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E402, E501 + compressed_tensors_moe_w4a16_flydsl, +) + +RoutingBuffers = tuple[ + torch.Tensor, # sorted_token_ids + torch.Tensor, # sorted_weights + torch.Tensor, # sorted_expert_ids + torch.Tensor, # num_valid_ids (shape [1], i32) + int, # sorted_size + int, # blocks +] + + +@pytest.mark.parametrize( + "num_tokens", [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384] +) +@pytest.mark.parametrize("inter_dim", [256, 512]) +def test_flydsl_moe(num_tokens: int, inter_dim: int): + device = "cuda" + topk = 8 + num_experts = 384 + hidden_size = 7168 + packed_factor = 8 + w13_num_shards = 2 + params_dtype = torch.bfloat16 + group_size = 32 + w2_scales_size = inter_dim + scale_factor = 0.01 + + num_groups_w2 = w2_scales_size // group_size + num_groups_w13 = hidden_size // group_size + + w13_weight = torch.randint( + 0, + 255, + (num_experts, hidden_size // packed_factor, w13_num_shards * inter_dim), + dtype=torch.int32, + device=device, + ) + + w2_weight = torch.randint( + 0, + 255, + (num_experts, inter_dim // packed_factor, hidden_size), + dtype=torch.int32, + device=device, + ) + w13_scale = scale_factor * torch.randn( + num_experts, + num_groups_w13, + w13_num_shards * inter_dim, + dtype=params_dtype, + device=device, + ) + w2_scale = scale_factor * torch.randn( + num_experts, num_groups_w2, hidden_size, dtype=params_dtype, device=device + ) + + w13_weight_packed = w13_weight.transpose(1, 2).contiguous().view(torch.uint8) + w2_weight_packed = w2_weight.transpose(1, 2).contiguous().view(torch.uint8) + w13_weight_scale = w13_scale.transpose(1, 2).contiguous() + w2_weight_scale = w2_scale.transpose(1, 2).contiguous() + + moe_quant_config = int4_w4a16_moe_quant_config( + w1_scale=w13_weight_scale, + w2_scale=w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) + score = torch.rand((num_tokens, num_experts), device=device, dtype=torch.float32) + topk_vals, topk_ids = torch.topk(score, k=topk, dim=1) + topk_weights = torch.softmax(topk_vals, dim=1).to(torch.float32) + x = torch.randn((num_tokens, hidden_size), dtype=torch.bfloat16, device=device) + out_ref = fused_experts( + x, + w13_weight_packed, + w2_weight_packed, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + apply_router_weight_on_input=False, + global_num_experts=num_experts, + expert_map=None, + quant_config=moe_quant_config, + ) + + w13 = w13_weight + w13 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w13) + w13 = w13.view(-1).contiguous() + + w2 = w2_weight + w2 = compressed_tensors_moe_w4a16_flydsl._gptq_int32_to_flydsl_packed(w2) + w2 = w2.view(-1).contiguous() + + w13_scale_flydsl = w13_scale + w2_scale_flydsl = w2_scale + + if group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1: + E, G, N = w13_scale.shape + w13_scale_flydsl = ( + w13_scale_flydsl.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1: + w13_scale_flydsl = w13_scale_flydsl.squeeze(1) + + if group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1: + E, G, N = w2_scale.shape + w2_scale_flydsl = ( + w2_scale_flydsl.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1: + w2_scale_flydsl = w2_scale_flydsl.squeeze(1) + + w13_scale_flydsl = w13_scale_flydsl.contiguous() + w2_scale_flydsl = w2_scale_flydsl.contiguous() + + w13.is_shuffled = True + w2.is_shuffled = True + + out = fused_flydsl_moe( + x, + w13, + w2, + num_experts, + inter_dim, + topk_weights, + topk_ids, + w1_scale=w13_scale_flydsl, + w2_scale=w2_scale_flydsl, + topk=topk_weights.shape[-1], + group_size=group_size, + doweight_stage1=False, + scale_is_bf16=True, + ) + + assert torch.allclose(out, out_ref, atol=0.5, rtol=0.1) + + +if __name__ == "__main__": + test_flydsl_moe(512, 256) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 7a393752f476..46dad3aa44b8 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -133,6 +133,7 @@ def with_default( "humming", "triton_unfused", "aiter", + "flydsl", "emulation", ] @@ -186,6 +187,7 @@ class KernelConfig: - "humming": Use Humming Mixed Precision kernels - "triton_unfused": Use Triton unfused MoE kernels - "aiter": Use AMD AITer kernels (ROCm only) + - "flydsl": Use AMD FlyDSL kernels (ROCm only) - "emulation": use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations. """ diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..bb1a95d3d519 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 256 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "256": { + "tile_m": 16, + "tile_n": 128, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..bb1a95d3d519 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 256 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "256": { + "tile_m": 16, + "tile_n": 128, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..bb1a95d3d519 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 256 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "256": { + "tile_m": 16, + "tile_n": 128, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..bb1a95d3d519 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 256 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 256 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "256": { + "tile_m": 16, + "tile_n": 128, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 256 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 256, + "tile_k2": 64 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..5bd96accd9c4 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "256": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 64 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 128 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..5bd96accd9c4 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "256": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 64 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 128 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..5bd96accd9c4 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "256": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 64 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 128 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json new file mode 100644 index 000000000000..5bd96accd9c4 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json @@ -0,0 +1,114 @@ +{ + "1": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 256, + "tile_k2": 128 + }, + "2": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "4": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 512, + "tile_n2": 128, + "tile_k2": 128 + }, + "8": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "16": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "24": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "32": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "48": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "64": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "128": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "256": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 256, + "tile_k2": 128 + }, + "512": { + "tile_m": 16, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 128 + }, + "1024": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 128, + "tile_n2": 128, + "tile_k2": 64 + }, + "2048": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "4096": { + "tile_m": 32, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 64 + }, + "8192": { + "tile_m": 64, + "tile_n": 64, + "tile_k": 64, + "tile_n2": 128, + "tile_k2": 128 + } +} \ No newline at end of file diff --git a/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py new file mode 100644 index 000000000000..cf49e01e6282 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused MoE Triton kernels.""" + +import functools +import json +import os + +import flydsl.compiler as flyc +import torch +from aiter.fused_moe import moe_sorting as aiter_moe_sorting +from aiter.ops.flydsl.kernels.moe_gemm_2stage import ( + compile_moe_gemm1, + compile_moe_gemm2, +) + +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.utils.torch_utils import direct_register_custom_op + +logger = init_logger(__name__) + +_FLYDSL_MOE_GEMM1_CACHE: dict = {} +_FLYDSL_MOE_GEMM2_CACHE: dict = {} + +_FLYDSL_MOE_DEFAULT_CONFIG = { + 1: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256}, + 2: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128}, + 4: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 128}, + 8: {"tile_m": 16, "tile_n": 64, "tile_k": 512, "tile_n2": 256, "tile_k2": 256}, + 16: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 256}, + 24: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 32: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 48: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 64: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 128, "tile_k2": 128}, + 128: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 256: {"tile_m": 16, "tile_n": 128, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 512: {"tile_m": 16, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 1024: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 2048: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64}, + 4096: {"tile_m": 32, "tile_n": 64, "tile_k": 128, "tile_n2": 256, "tile_k2": 256}, + 8192: {"tile_m": 64, "tile_n": 64, "tile_k": 64, "tile_n2": 256, "tile_k2": 64}, +} + + +def moe_sorting( + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + *, + num_experts: int, + model_dim: int, + block_m: int, +): + topk_ids_i32 = topk_ids.to(torch.int32) + topk_w_f32 = topk_weights.to(torch.float32) + sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids, _moe_buf = ( + aiter_moe_sorting( + topk_ids_i32, + topk_w_f32, + num_experts, + model_dim, + torch.float16, + block_m, + ) + ) + if num_valid_ids.numel() > 1: + num_valid_ids = num_valid_ids[:1].contiguous() + return sorted_ids, sorted_w, sorted_expert_ids, num_valid_ids + + +def build_routing_buffers( + *, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + num_experts: int, + model_dim: int, + tile_m: int, +): + res = moe_sorting( + topk_ids, + topk_weights, + num_experts=num_experts, + model_dim=model_dim, + block_m=tile_m, + ) + if res is None: + raise RuntimeError( + "aiter moe_sorting failed/unavailable; cannot build routing buffers." + ) + sorted_token_ids, sorted_weights, sorted_expert_ids, num_valid_ids = res + + sorted_token_ids = sorted_token_ids.contiguous() + sorted_weights = sorted_weights.contiguous() + sorted_expert_ids = sorted_expert_ids.contiguous() + sorted_size = int(sorted_token_ids.numel()) + blocks = int(sorted_expert_ids.numel()) + return ( + sorted_token_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + sorted_size, + blocks, + ) + + +@functools.lru_cache +def try_get_optimal_config(num_experts, inter_dim): + device_name = current_platform.get_device_name().replace(" ", "_") + json_file_name = ( + f"E={num_experts},N={inter_dim},device_name={device_name}," + "dtype=int4_w4a16,backend=flydsl.json" + ) + config_file_path = os.path.join( + os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name + ) + if os.path.exists(config_file_path): + with open(config_file_path) as f: + logger.info_once( + "Using tuned FlyDSL MoE config from %s", + config_file_path, + scope="global", + ) + tuned_config = json.load(f) + return {int(key): val for key, val in tuned_config.items()} + + logger.warning_once( + "Using default FlyDSL MoE config. Performance might be sub-optimal! " + "Config file not found at %s", + config_file_path, + scope="local", + ) + return _FLYDSL_MOE_DEFAULT_CONFIG + + +def fused_flydsl_moe_impl( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + num_experts: int, + inter_dim: int, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, + topk: int = 8, + group_size: int = 32, + doweight_stage1: bool = False, + in_dtype: str = "int4_bf16", + out_dtype: str = "bf16", + scale_is_bf16: bool = True, + tile_m: int | None = None, + tile_n: int | None = None, + tile_k: int | None = None, + tile_n2: int | None = None, + tile_k2: int | None = None, +) -> torch.Tensor: + device = hidden_states.device + tokens = hidden_states.shape[0] + model_dim = hidden_states.shape[1] + + tuned_config = {} + if tile_m and tile_n and tile_k and tile_n2 and tile_k2: + tuned_config["tile_m"] = tile_m + tuned_config["tile_n"] = tile_n + tuned_config["tile_k"] = tile_k + tuned_config["tile_n2"] = tile_n2 + tuned_config["tile_k2"] = tile_k2 + else: + tuned_config = try_get_optimal_config(num_experts, inter_dim) + tuned_config = tuned_config[ + min(tuned_config.keys(), key=lambda x: abs(x - tokens)) + ] + out_torch_dtype = torch.bfloat16 if out_dtype == "bf16" else torch.float16 + + tile_m = tuned_config["tile_m"] + tile_n = tuned_config["tile_n"] + tile_k = tuned_config["tile_k"] + tile_n2 = tuned_config["tile_n2"] + tile_k2 = tuned_config["tile_k2"] + + routing = build_routing_buffers( + topk_ids=topk_ids, + topk_weights=topk_weights, + num_experts=num_experts, + model_dim=model_dim, + tile_m=tile_m, + ) + ( + sorted_token_ids, + sorted_weights, + sorted_expert_ids, + num_valid_ids, + sorted_size, + blocks, + ) = routing + + scale_x_1d = torch.empty((0,), device=device, dtype=torch.float32) + sorted_weights_1d = sorted_weights.view(-1).contiguous() + out_stage1 = torch.empty( + (tokens, topk, inter_dim), device=device, dtype=out_torch_dtype + ) + + stream = torch.cuda.current_stream() + + key1 = ( + model_dim, + inter_dim, + num_experts, + topk, + in_dtype, + out_dtype, + group_size, + tile_m, + tile_n, + tile_k, + bool(doweight_stage1), + False, + ) + + compiled_exe1 = _FLYDSL_MOE_GEMM1_CACHE.get(key1) + if compiled_exe1 is None: + exe1 = compile_moe_gemm1( + model_dim=model_dim, + inter_dim=inter_dim, + experts=num_experts, + topk=topk, + in_dtype=in_dtype, + out_dtype=out_dtype, + group_size=group_size, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + doweight_stage1=bool(doweight_stage1), + use_cshuffle_epilog=False, + scale_is_bf16=scale_is_bf16, + ) + compiled_exe1 = flyc.compile( + exe1, + out_stage1, + hidden_states, + w1, + scale_x_1d, + w1_scale, + sorted_token_ids, + sorted_expert_ids, + sorted_weights_1d, + num_valid_ids, + tokens, + inter_dim, + model_dim, + int(blocks), + stream, + ) + _FLYDSL_MOE_GEMM1_CACHE[key1] = compiled_exe1 + + compiled_exe1( + out_stage1, + hidden_states, + w1, + scale_x_1d, + w1_scale, + sorted_token_ids, + sorted_expert_ids, + sorted_weights_1d, + num_valid_ids, + tokens, + inter_dim, + model_dim, + int(blocks), + stream, + ) + + a2_1d = out_stage1.view(-1).contiguous() + a2_scale_1d = torch.empty((0,), device=device, dtype=torch.float32) + out_stage2 = torch.empty((tokens, model_dim), device=device, dtype=out_torch_dtype) + doweight_stage2 = not bool(doweight_stage1) + + key2 = ( + model_dim, + inter_dim, + num_experts, + topk, + in_dtype, + out_dtype, + group_size, + tile_m, + tile_n2, + tile_k2, + bool(doweight_stage2), + ) + + compiled_exe2 = _FLYDSL_MOE_GEMM2_CACHE.get(key2) + if compiled_exe2 is None: + exe2 = compile_moe_gemm2( + model_dim=model_dim, + inter_dim=inter_dim, + experts=num_experts, + topk=topk, + in_dtype=in_dtype, + out_dtype=out_dtype, + group_size=group_size, + tile_m=tile_m, + tile_n=tile_n2, + tile_k=tile_k2, + doweight_stage2=bool(doweight_stage2), + scale_is_bf16=scale_is_bf16, + ) + compiled_exe2 = flyc.compile( + exe2, + out_stage2, + a2_1d, + w2, + a2_scale_1d, + w2_scale, + sorted_token_ids, + sorted_expert_ids, + sorted_weights_1d, + num_valid_ids, + tokens, + model_dim, + inter_dim, + int(blocks), + stream, + ) + _FLYDSL_MOE_GEMM2_CACHE[key2] = compiled_exe2 + + out_stage2.zero_() + compiled_exe2( + out_stage2, + a2_1d, + w2, + a2_scale_1d, + w2_scale, + sorted_token_ids, + sorted_expert_ids, + sorted_weights_1d, + num_valid_ids, + tokens, + model_dim, + inter_dim, + int(blocks), + stream, + ) + return out_stage2 + + +def fused_flydsl_moe_impl_fake( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + num_experts: int, + inter_dim: int, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, + topk: int = 8, + group_size: int = 32, + doweight_stage1: bool = False, + in_dtype: str = "int4_bf16", + out_dtype: str = "bf16", + scale_is_bf16: bool = True, + tile_m: int | None = None, + tile_n: int | None = None, + tile_k: int | None = None, + tile_n2: int | None = None, + tile_k2: int | None = None, +) -> torch.Tensor: + return torch.empty_like(hidden_states) + + +direct_register_custom_op( + op_name="fused_flydsl_moe_impl", + op_func=fused_flydsl_moe_impl, + fake_impl=fused_flydsl_moe_impl_fake, +) + + +def fused_flydsl_moe( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + num_experts: int, + inter_dim: int, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + w1_scale: torch.Tensor | None = None, + w2_scale: torch.Tensor | None = None, + topk: int = 8, + group_size: int = 32, + doweight_stage1: bool = False, + in_dtype: str = "int4_bf16", + out_dtype: str = "bf16", + scale_is_bf16: bool = True, + config: dict | None = None, +) -> torch.Tensor: + tile_m = None + tile_n = None + tile_k = None + tile_n2 = None + tile_k2 = None + if config is not None: + tile_m = config.get("tile_m") + tile_n = config.get("tile_n") + tile_k = config.get("tile_k") + tile_n2 = config.get("tile_n2") + tile_k2 = config.get("tile_k2") + return torch.ops.vllm.fused_flydsl_moe_impl( + hidden_states=hidden_states, + w1=w1, + w2=w2, + num_experts=num_experts, + inter_dim=inter_dim, + topk_weights=topk_weights, + topk_ids=topk_ids, + w1_scale=w1_scale, + w2_scale=w2_scale, + topk=topk, + group_size=group_size, + doweight_stage1=doweight_stage1, + in_dtype=in_dtype, + out_dtype=out_dtype, + scale_is_bf16=scale_is_bf16, + tile_m=tile_m, + tile_n=tile_n, + tile_k=tile_k, + tile_n2=tile_n2, + tile_k2=tile_k2, + ) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 69c27551bf1a..9a75d6a3f1a0 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -197,6 +197,7 @@ def _needs_intermediate_size_param(self, quant_method: FusedMoEMethodBase) -> bo "AutoGPTQMoEMethod", "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod", + "CompressedTensorsW4A16FlydslMoEMethod", ) def _ensure_moe_quant_config_init(self): @@ -610,6 +611,7 @@ def weight_loader( "CompressedTensorsWNA16MarlinMoEMethod", "CompressedTensorsWNA16MoEMethod", "CompressedTensorsWNA16RDNA3MoEMethod", + "CompressedTensorsW4A16FlydslMoEMethod", ): if is_transposed: loaded_weight = loaded_weight.t().contiguous() diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 0c3a434ba5f7..2e45e0f298ba 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -7,8 +7,10 @@ from compressed_tensors.quantization import ( ActivationOrdering, QuantizationStrategy, + QuantizationType, ) +from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoEMethodBase, @@ -115,7 +117,28 @@ def get_moe_method( return rocm_moe_rdna.make_method( weight_quant, input_quant, layer.moe_config ) + from vllm.platforms.rocm import on_gfx950 + + vllm_config = get_current_vllm_config() + is_lora_disabled = vllm_config.lora_config is None + moe_backend = vllm_config.kernel_config.moe_backend + if ( + weight_quant.strategy == QuantizationStrategy.GROUP + and weight_quant.type == QuantizationType.INT + and group_size == 32 + and weight_quant.num_bits == 4 + and is_lora_disabled + and on_gfx950() + and moe_backend == "flydsl" + ): + from .compressed_tensors_moe_w4a16_flydsl import ( + CompressedTensorsW4A16FlydslMoEMethod, + ) + logger.info_once("Using CompressedTensorsW4A16FlydslMoEMethod") + return CompressedTensorsW4A16FlydslMoEMethod( + weight_quant, input_quant, layer.moe_config + ) from .compressed_tensors_moe_wna16 import ( CompressedTensorsWNA16MoEMethod, ) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py new file mode 100644 index 000000000000..f8faddbd07bf --- /dev/null +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +from aiter.ops.shuffle import shuffle_weight +from compressed_tensors.quantization import ( + QuantizationArgs, +) + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import ( + RoutedExperts, + SharedExperts, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa E501 + CompressedTensorsMoEMethod, +) +from vllm.model_executor.utils import set_weight_attrs + +logger = init_logger(__name__) + + +def _pack_shuffled_int8_to_packed_int4_no_perm(x_shuf_i8: torch.Tensor) -> torch.Tensor: + """Pack a preshuffled int8 tensor (values in [-8, 7]) into packed int4 bytes. + Each contiguous 8-value block [v0..v7] -> 4 bytes: + b0=(v4<<4)|v0, b1=(v5<<4)|v1, b2=(v6<<4)|v2, b3=(v7<<4)|v3. + This matches the 7-op in-kernel unpack sequence and avoids any v_perm. + """ + flat = x_shuf_i8.contiguous().view(-1).to(torch.int16) + assert flat.numel() % 8 == 0 + u = (flat & 0xF).to(torch.uint8).view(-1, 8) + out = torch.empty((u.shape[0], 4), device=u.device, dtype=torch.uint8) + out[:, 0] = u[:, 0] | (u[:, 4] << 4) + out[:, 1] = u[:, 1] | (u[:, 5] << 4) + out[:, 2] = u[:, 2] | (u[:, 6] << 4) + out[:, 3] = u[:, 3] | (u[:, 7] << 4) + return out.view(-1).to(torch.int8) + + +def _unpack_gptq_int32_to_signed_int4(w_int32): + """Unpack GPTQ int32 [E, K//8, N] to signed int4 values [E, N, K] (as int8). + Shared by both the packed-int4 and bf16-dequant paths. + """ + E = w_int32.shape[0] + # [E, K//8, N] -> transpose -> [E, N, K//8] + w = w_int32.transpose(1, 2).contiguous() + N = w.shape[1] + K_div8 = w.shape[2] + K = K_div8 * 8 + + # Unpack int32 -> 8 x uint4 values along K + w_expanded = w.unsqueeze(-1).expand(E, N, K_div8, 8) # [E, N, K//8, 8] + shifts = torch.arange(8, device=w.device) * 4 # [0, 4, 8, ..., 28] + nibbles = ((w_expanded >> shifts) & 0xF).to(torch.int8) # [E, N, K//8, 8] + nibbles = nibbles.reshape(E, N, K) # [E, N, K] unsigned int4 as int8 + + # Convert unsigned [0,15] to signed [-8,7] + signed = nibbles.to(torch.int16) - 8 + signed = signed.to(torch.int8) # [E, N, K] signed int4 as int8 + return signed + + +def _gptq_int32_to_flydsl_packed(w_int32): + """Convert GPTQ int32 [E, K//8, N] to FlyDSL shuffled packed int4 [E, N, K//2]. + Steps: + 1. Unpack int32 to individual signed int4 values (as int8) + 2. Apply FlyDSL preshuffle (on individual int8 values) + 3. Pack with FlyDSL's interleaved int4 packing + """ + signed = _unpack_gptq_int32_to_signed_int4(w_int32) + E, N, K = signed.shape + + # FlyDSL preshuffle (operates on individual values) + shuffled = shuffle_weight(signed, layout=(16, 16)) + + # FlyDSL interleaved int4 packing + packed = _pack_shuffled_int8_to_packed_int4_no_perm(shuffled).contiguous() + return packed.view(E, N, K // 2) + + +class CompressedTensorsW4A16FlydslMoEMethod(CompressedTensorsMoEMethod): + def __init__( + self, + weight_quant: QuantizationArgs, + input_quant: QuantizationArgs | None, + moe: FusedMoEConfig, + layer_name: str | None = None, + ): + super().__init__(moe) + self.weight_quant = weight_quant + self.input_quant = input_quant + # Extract properties from weight_quant + assert weight_quant.num_bits == 4 + self.num_bits = weight_quant.num_bits + self.packed_factor = 32 // weight_quant.num_bits + self.strategy = weight_quant.strategy + # channelwise is not supported by this kernel + assert weight_quant.strategy == "group" + assert weight_quant.group_size == 32 + self.group_size = weight_quant.group_size + # grouped actorder isn't supported by this kernel + assert weight_quant.actorder != "group" + assert weight_quant.symmetric, ( + "Only symmetric quantization is supported for MoE" + ) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + self.num_experts = num_experts + self.inter_dim = intermediate_size_per_partition + # Will transpose the loaded weight along the + # intermediate and hidden dim sizes. Will + # shard for TP along the transposed dims + extra_weight_attrs.update( + {"is_transposed": True, "quant_method": self.strategy} + ) + w13_num_shards = 2 if self.moe.is_act_and_mul else 1 + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size // self.packed_factor, + w13_num_shards * intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition // self.packed_factor, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_packed", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + w2_scales_size = intermediate_size_per_partition + + if self.strategy == "channel": + num_groups_w2 = num_groups_w13 = 1 + self.group_size = -1 + else: + num_groups_w2 = w2_scales_size // self.group_size + num_groups_w13 = hidden_size // self.group_size + + w13_scale = torch.nn.Parameter( + torch.ones( + num_experts, + num_groups_w13, + w13_num_shards * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_scale) + set_weight_attrs(w13_scale, extra_weight_attrs) + + w2_scale = torch.nn.Parameter( + torch.ones(num_experts, num_groups_w2, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_scale) + set_weight_attrs(w2_scale, extra_weight_attrs) + set_weight_attrs(w2_scale, {"load_full_w2": False}) + + w2_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + layer.register_parameter("w2_weight_shape", w2_weight_shape) + set_weight_attrs(w2_weight_shape, extra_weight_attrs) + w13_weight_shape = torch.nn.Parameter( + torch.empty(num_experts, 2), requires_grad=False + ) + + layer.register_parameter("w13_weight_shape", w13_weight_shape) + set_weight_attrs(w13_weight_shape, extra_weight_attrs) + + w13_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_g_idx", w13_g_idx) + set_weight_attrs(w13_g_idx, extra_weight_attrs) + + w2_g_idx = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_g_idx", w2_g_idx) + set_weight_attrs(w2_g_idx, extra_weight_attrs) + + w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices) + set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs) + + w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) + set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + + layer.a13_scale = None + layer.a2_scale = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Reconfigure packed weights and scales to match flydsl_w4a16 format + + # Convert w13 weights + w13 = layer.w13_weight_packed.data + w13 = _gptq_int32_to_flydsl_packed(w13) + w13 = w13.view(-1).contiguous() + layer.w13_weight_packed = torch.nn.Parameter(w13, requires_grad=False) + + # Convert w2 weights + w2 = layer.w2_weight_packed.data + w2 = _gptq_int32_to_flydsl_packed(w2) + w2 = w2.view(-1).contiguous() + layer.w2_weight_packed = torch.nn.Parameter(w2, requires_grad=False) + + # Convert scales for FlyDSL: + # per-row: [E, 1, N] -> squeeze -> [E, N] + # groupwise: [E, K//gs, N] -> keep as-is (Opt 0: cache-friendly layout) + w13_scale = layer.w13_weight_scale.data + if self.group_size > 0 and w13_scale.dim() == 3 and w13_scale.shape[1] > 1: + E, G, N = w13_scale.shape + w13_scale = ( + w13_scale.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w13_scale.dim() == 3 and w13_scale.shape[1] == 1: + # Per-row: squeeze [E, 1, N] -> [E, N] + w13_scale = w13_scale.squeeze(1) + layer.w13_weight_scale = torch.nn.Parameter( + w13_scale.contiguous(), requires_grad=False + ) + + w2_scale = layer.w2_weight_scale.data + if self.group_size > 0 and w2_scale.dim() == 3 and w2_scale.shape[1] > 1: + E, G, N = w2_scale.shape + w2_scale = ( + w2_scale.view(E, G // 2, 2, N) + .permute(0, 1, 3, 2) + .contiguous() + .view(-1) + .contiguous() + ) + elif w2_scale.dim() == 3 and w2_scale.shape[1] == 1: + # Per-row: squeeze [E, 1, N] -> [E, N] + w2_scale = w2_scale.squeeze(1) + layer.w2_weight_scale = torch.nn.Parameter( + w2_scale.contiguous(), requires_grad=False + ) + + layer.w13_weight_packed.is_shuffled = True + layer.w2_weight_packed.is_shuffled = True + layer.is_aiter_converted = True + + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + assert self.num_bits == 4 + return int4_w4a16_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, self.group_size], + ) + + def select_gemm_impl( + self, + prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, + layer: torch.nn.Module, + ) -> mk.FusedMoEExpertsModular: + raise NotImplementedError + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + from vllm.model_executor.layers.fused_moe.fused_flydsl_moe import ( + fused_flydsl_moe, + ) + + assert self.moe_quant_config is not None + + return fused_flydsl_moe( + x, + layer.w13_weight_packed, + layer.w2_weight_packed, + self.num_experts, + self.inter_dim, + topk_weights, + topk_ids, + w1_scale=self.moe_quant_config.w1_scale, + w2_scale=self.moe_quant_config.w2_scale, + topk=topk_weights.shape[-1], + group_size=self.group_size, + doweight_stage1=layer.apply_router_weight_on_input, + scale_is_bf16=True, + )