Skip to content

[HIP] [CK] Revert " Re-submit Relu2 Tier 2 MoE GEMM support (#6376)" #6463

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zufayu wants to merge 1 commit into from
Closed

[HIP] [CK] Revert " Re-submit Relu2 Tier 2 MoE GEMM support (#6376)" #6463

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 4 additions & 84 deletions aiter/fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -799,8 +799,6 @@ def resolve_activation_dtype(
q_dtype_a = _bound_split(
M, _SWIGLU_MXFP4_BF16_BOUND, dtypes.bf16, dtypes.fp4x2
)
elif activation == ActivationType.Relu2:
q_dtype_a = dtypes.bf16
elif activation == ActivationType.Swiglu or gate_mode == GateMode.INTERLEAVE:
if gfx != "gfx950":
q_dtype_a = dtypes.bf16
Expand Down Expand Up @@ -2985,31 +2983,6 @@ def get_2stage_cfgs(
doweight_stage1=doweight_stage1,
gate_mode=gate_mode,
)
if activation == ActivationType.Relu2:
if use_g1u1:
raise NotImplementedError(
"ActivationType.Relu2 is gate-only (non-gated) and does not "
"support use_g1u1=True; pass use_g1u1=False."
)
if doweight_stage1:
raise NotImplementedError(
"ActivationType.Relu2 CK-Tile stage1 instances are generated "
"with MulRoutedWeight=False; doweight_stage1=True is not "
"supported and would silently drop routed weights."
)
if has_stage1_bias or has_stage2_bias:
raise NotImplementedError(
"ActivationType.Relu2 CK-Tile path does not support "
"per-expert bias (has_stage1_bias/has_stage2_bias); the "
"returned metadata hardcodes has_bias=False."
)
if gate_mode != GateMode.SEPARATED:
raise NotImplementedError(
"ActivationType.Relu2 CK-Tile stage1 instance "
"(kFFN_gemm1_gate_only) only supports the gate-only/"
f"separated weight layout; gate_mode={gate_mode.value!r} "
"is not supported."
)
# Configs are keyed on (gfx, cu_num, ...) so archs that share a cu_num
# (e.g. gfx950 vs gfx1250, both report 256 CU) don't collide. Legacy CSVs
# without a `gfx` column are backfilled from cu_num at load time via
Expand Down Expand Up @@ -3235,18 +3208,6 @@ def _lookup_cfg(c2s):
f"[fused_moe] discarding Opus tuned config for unsupported "
f"activation {activation}; using default heuristics"
)
elif (
kn1.startswith("flydsl_") or kn2.startswith("flydsl_")
) and activation not in (
ActivationType.Silu,
ActivationType.Swiglu,
ActivationType.Situv2,
):
cfg = None
logger.warning(
f"[fused_moe] discarding FlyDSL tuned config for unsupported "
f"activation {activation}; using default heuristics"
)
elif _disable_inline_sort and _is_inline_sort_cfg(kn1, kn2):
cfg = None
logger.warning("[fused_moe] discarding tuned inline-sort config")
Expand Down Expand Up @@ -3582,9 +3543,7 @@ def get_block_m() -> int:
is_opus2 = _opus_a8w4.is_opus_a8w4_stage2_kernel(kernelName2)
if is_opus1 or is_flydsl1 or is_flydsl2:
enable_bias = (
activation == ActivationType.Swiglu
and _needs_swiglu_bias_support(dtype, q_type)
and q_dtype_w == dtypes.fp4x2
_needs_swiglu_bias_support(dtype, q_type) and q_dtype_w == dtypes.fp4x2
)
_s1_fp4q = is_flydsl1 and "_fp4" in kernelName1.split("_t")[-1]
if is_opus1:
Expand Down Expand Up @@ -3723,40 +3682,6 @@ def get_block_m() -> int:
and is_shuffled
and cktile_mxfp4_ok
)
if (
activation == ActivationType.Relu2
and not use_g1u1
and q_type == QuantType.per_1x32
and q_dtype_w == dtypes.fp4x2
and q_dtype_a == dtypes.bf16
and dtype == dtypes.bf16
and not has_stage2_scatter
and not has_activation_scales
and bool(opus_weights_shuffled)
and get_gfx() == "gfx950"
):
_cktile_block_m = 16 if token < 2048 else 32 if token < 16384 else 64
return MOEMetadata(
functools.partial(
cktile_moe_stage1,
n_pad_zeros=intermediate_pad // 64 * 64,
k_pad_zeros=hidden_pad // 128 * 128,
activation=activation,
split_k=1,
dtype=dtype,
),
functools.partial(
cktile_moe_stage2,
n_pad_zeros=hidden_pad // 64 * 64,
k_pad_zeros=intermediate_pad // 128 * 128,
activation=activation,
),
_cktile_block_m,
1,
run_1stage,
has_bias=False,
stage2_has_bias=False,
)
if q_type == QuantType.per_1x32 and q_dtype_w == dtypes.i4x2:
# Untuned a16wi4 fallback: one shape-safe config on the shared a16w-mix port.
# Tiles belong in the tuned CSV, not in a heuristic here. ksplit is 0 because
Expand Down Expand Up @@ -3928,10 +3853,7 @@ def get_block_m() -> int:
f"[fused_moe] no tuned FlyDSL config for {keys}, "
f"using heuristic FlyDSL fallback ({kn1=}, {kn2=})"
)
enable_bias = (
activation == ActivationType.Swiglu
and _needs_swiglu_bias_support(dtype, q_type)
)
enable_bias = _needs_swiglu_bias_support(dtype, q_type)
return MOEMetadata(
functools.partial(
_flydsl_stage1_wrapper,
Expand Down Expand Up @@ -4251,8 +4173,7 @@ def fused_moe_2stages(
and (
q_dtype_a in [dtypes.bf16, dtypes.fp16]
and (
activation
in (ActivationType.Swiglu, ActivationType.Situv2, ActivationType.Relu2)
activation in (ActivationType.Swiglu, ActivationType.Situv2)
or gate_mode == GateMode.INTERLEAVE
)
or (q_dtype_a in [dtypes.fp4x2] and metadata.ksplit > 1 and is_shuffled)
Expand Down Expand Up @@ -4493,8 +4414,7 @@ def fused_moe_2stages(
and w1.dtype == dtypes.fp4x2
and (
q_dtype_a in [dtypes.bf16, dtypes.fp16]
and activation
in (ActivationType.Swiglu, ActivationType.Situv2, ActivationType.Relu2)
and activation in (ActivationType.Swiglu, ActivationType.Situv2)
or (metadata.ksplit > 1 and is_shuffled)
)
):
Expand Down
1 change: 0 additions & 1 deletion aiter/ops/quant.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -709,7 +709,6 @@ def get_torch_act(aType):
ActivationType.Silu: F.silu,
ActivationType.Gelu: F.gelu,
ActivationType.GeluTanh: lambda x: F.gelu(x, approximate="tanh"),
ActivationType.Relu2: lambda x: F.relu(x) ** 2,
}
return tmp.get(aType, NotImplementedError)

Expand Down
13 changes: 4 additions & 9 deletions csrc/ck_tile_gemm_moe_2stages/gen_instances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ def gen_instance(self, k: kernelInstance):
int KBatch = k_batch.has_value() ? k_batch.value() : 1;
int stride_A = K;
int stride_B = K;
int stride_C = KBatch > 1 ? N : N / {2 if (k.stage == 1 and k.ActOP != "relu2") else 1}; //gemm1 gate+up need / 2; relu2 is gate-only (non-gated), needs full-width N.
int stride_C = KBatch > 1 ? N : N / {3 - k.stage}; //gemm1 gate+up need / 2.
void *sorted_weights_ptr = topk_weight.has_value() ? topk_weight.value().data_ptr() : nullptr;

{{INSTANCE_CONTENT}}
Expand Down Expand Up @@ -190,7 +190,6 @@ def gen_instance(self, k: kernelInstance):
ck_tile::tuple<>,
row_major,
{"ck_tile::MoeFlatmmKind::kFFN_gemm1_split_k" if self.is_split_k else
"ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_only" if (k.stage == 1 and self.activation == act_dict["relu2"]) else
"ck_tile::MoeFlatmmKind::kFFN_gemm1_gate_up" if k.stage == 1 else
"ck_tile::MoeFlatmmKind::kFFN_gemm2"},
ck_tile::element_wise::PassThrough,
Expand Down Expand Up @@ -284,9 +283,7 @@ def validate_and_format(template: str, mapping: dict) -> str:
"(acc_data_type)": dtype_dict[self.acc_dtype],
"(c_data_type)": dtype_dict[self.c_dtype],
"(activation)": self.activation,
"(has_bias)": (
"true" if self.activation == 2 and not self.is_split_k else "false"
),
"(has_bias)": "true" if self.activation == 2 else "false",
"(split_k)": "true" if self.is_split_k else "false",
}
format_args = {str(key): value.name for key, value in mapping.items()}
Expand Down Expand Up @@ -622,7 +619,7 @@ def generate_common_header(working_path, dispatch_files, manifest_files):
quant_type = "1x32"

acc_type = "float"
act_types = ["silu", "swiglu", "relu2"]
act_types = ["silu", "swiglu"]
c_dtypes = ["bf16"]
is_split_k_l = [True, False]

Expand Down Expand Up @@ -653,13 +650,11 @@ def generate_common_header(working_path, dispatch_files, manifest_files):
for a_type, c_dtype, act_type, is_split_k in itertools.product(
a_types, c_dtypes, act_types, is_split_k_l
):
has_bias = act_type == "swiglu" and not is_split_k
has_bias = act_type == "swiglu"

# a8w8 do not support
if a_type in ["fp8", "bf8"] and is_split_k:
continue
if act_type == "relu2" and is_split_k:
continue
codegen = cktile_moe_2stage_gemm_codegen(
args.working_path,
a_type,
Expand Down
6 changes: 1 addition & 5 deletions csrc/ck_tile_gemm_moe_2stages/include/moe_cktile2stages_common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,11 +235,7 @@ void moe_gemm(const MoeFlatmmHostArgs& args, const ck_stream_config& s)

// TODO: support more act type.
using FusedAct =
std::conditional_t<ActivationOp == 2,
ck_tile::moe::Swiglu,
std::conditional_t<ActivationOp == 3,
ck_tile::moe::MoeRelu2,
ck_tile::moe::MoeSilu>>;
std::conditional_t<ActivationOp == 2, ck_tile::moe::Swiglu, ck_tile::moe::MoeSilu>;

using Kernel = ck_tile::MoeFlatmmKernel<TilePartitioner,
CodegenFlatmmPipeline,
Expand Down
28 changes: 0 additions & 28 deletions csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
#include "moe_cktile2stages_name_dispatch.h"
#include "py_itfs_common.h"
#include "moe_cktile2stages_heuristic_dispatch_common.h"
#include "aiter_enum.h"
#include <cmath>

template <typename ADataType,
Expand Down Expand Up @@ -243,34 +242,7 @@ MoeKernel moe_dispatch(int M, int N, int K, int block_m, int activation, bool ha
false>::dispatch(M, N, K, block_m);
}
}
else if(activation == static_cast<int>(ActivationType::Relu2) && !has_bias && split_k == 1)
{
if(stage == 1)
{
return moe_gemm1_heuristic_dispatcher<ADataType,
BDataType,
AccDataType,
CDataType,
3,
false,
false>::dispatch(M, N, K, block_m);
}
else
{
return moe_gemm2_heuristic_dispatcher<ADataType,
BDataType,
AccDataType,
CDataType,
3,
false,
false>::dispatch(M, N, K, block_m);
}
}
}
TORCH_CHECK(false,
"moe_dispatch: unsupported (activation=", activation,
", has_bias=", has_bias, ", split_k=", split_k,
") combination \u2014 no generated kernel instance covers this case");
}

torch::Tensor cktile_moe_gemm1(torch::Tensor& XQ,
Expand Down
1 change: 0 additions & 1 deletion csrc/ck_tile_gemm_moe_2stages/moe_cktile2stages_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
"silu": 0,
# "gelu": 1,
"swiglu": 2,
"relu2": 3,
}


Expand Down
1 change: 0 additions & 1 deletion csrc/include/aiter_enum.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@ enum class ActivationType : int
Swiglu = 2,
Situv2 = 3,
GeluTanh = 4,
Relu2 = 5,
};

enum class QuantType : int
Expand Down
1 change: 0 additions & 1 deletion csrc/include/rocm_ops.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,6 @@ namespace py = pybind11;
.value("Swiglu", ActivationType::Swiglu) \
.value("Situv2", ActivationType::Situv2) \
.value("GeluTanh", ActivationType::GeluTanh) \
.value("Relu2", ActivationType::Relu2) \
.export_values(); \
pybind11::enum_<MlaVersion>(m, "MlaVersion") \
.value("V32", MlaVersion::V32) \
Expand Down
Loading