iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][Perf] Enable W4A16 FlyDSL MoE (#44400)

Browse Source 
Signed-off-by: amd-asalykov <asalykov@amd.com>
Signed-off-by: Amanzhol Salykov <asalykov@amd.com>
This commit is contained in:
Amanzhol Salykov
2026-06-14 00:14:39 -07:00
committed by GitHub
parent 9548a1887f
commit 725c3bc808
15 changed files with 2173 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/kernels/benchmark_flydsl_moe_w4a16.py
+277
View File
@@ -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")
tests/kernels/moe/test_flydsl_moe.py
+179
View File
@@ -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)
vllm/config/kernel.py
+2
View File
@@ -133,6 +133,7 @@ MoEBackend = Literal[
"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.
"""
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=256,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350X,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI350_OAM,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/configs/E=384,N=512,device_name=AMD_Instinct_MI355_OAM,dtype=int4_w4a16,backend=flydsl.json
+114
View File
@@ -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
}
}
vllm/model_executor/layers/fused_moe/fused_flydsl_moe.py
+430
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/routed_experts.py
+2
View File
@@ -197,6 +197,7 @@ class RoutedExperts(PluggableLayer):
"AutoGPTQMoEMethod",
"CompressedTensorsWNA16MarlinMoEMethod",
"CompressedTensorsWNA16MoEMethod",
"CompressedTensorsW4A16FlydslMoEMethod",
)
def _ensure_moe_quant_config_init(self):
@@ -610,6 +611,7 @@ class RoutedExperts(PluggableLayer):
"CompressedTensorsWNA16MarlinMoEMethod",
"CompressedTensorsWNA16MoEMethod",
"CompressedTensorsWNA16RDNA3MoEMethod",
"CompressedTensorsW4A16FlydslMoEMethod",
):
if is_transposed:
loaded_weight = loaded_weight.t().contiguous()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+23
View File
@@ -7,8 +7,10 @@ from compressed_tensors import CompressionFormat
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 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
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,
)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a16_flydsl.py
+348
View File
@@ -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,
)