217 lines
7.5 KiB
Diff
217 lines
7.5 KiB
Diff
diff --git a/frontier/profiling/moe/moe_vllm_kernel.py b/frontier/profiling/moe/moe_vllm_kernel.py
|
|
--- a/frontier/profiling/moe/moe_vllm_kernel.py
|
|
+++ b/frontier/profiling/moe/moe_vllm_kernel.py
|
|
@@ -34,6 +34,7 @@ try:
|
|
import vllm
|
|
VLLM_VERSION = vllm.__version__
|
|
|
|
+ from vllm import _custom_ops as ops
|
|
# Import vLLM 0.10.x functions
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import (
|
|
fused_moe_kernel,
|
|
@@ -232,7 +233,7 @@ def _invoke_kernel(
|
|
"""
|
|
# Determine compute_type - for FP8, we accumulate in FP16/BF16
|
|
if use_fp8:
|
|
- compute_type = tl.float16 # FP8 accumulates in FP16
|
|
+ compute_type = tl.bfloat16
|
|
else:
|
|
dtype = A.dtype
|
|
if dtype == torch.bfloat16:
|
|
@@ -275,7 +276,9 @@ def _run_fused_moe_iteration(
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
intermediate_cache1: torch.Tensor,
|
|
intermediate_cache2: torch.Tensor,
|
|
+ intermediate_cache3: torch.Tensor,
|
|
+ output: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
sorted_token_ids: torch.Tensor,
|
|
expert_ids: torch.Tensor,
|
|
@@ -292,8 +295,17 @@ def _run_fused_moe_iteration(
|
|
per_channel_quant: bool = False,
|
|
block_shape: Optional[List[int]] = None,
|
|
) -> None:
|
|
+ first_input = A
|
|
+ first_A_scale = A_scale
|
|
+ if use_fp8:
|
|
+ group_size = block_dims[1] if block_dims else 128
|
|
+ first_input, first_A_scale = quantize_activations_to_fp8(
|
|
+ A,
|
|
+ group_size=group_size,
|
|
+ )
|
|
+
|
|
_invoke_kernel(
|
|
- A=A.contiguous(),
|
|
+ A=first_input.contiguous(),
|
|
B=w1.contiguous(),
|
|
C=intermediate_cache1.contiguous(),
|
|
topk_weights=topk_weights.contiguous(),
|
|
@@ -305,15 +316,18 @@ def _run_fused_moe_iteration(
|
|
mul_routed_weight=False,
|
|
top_k=top_k,
|
|
config=config,
|
|
- A_scale=A_scale,
|
|
+ A_scale=first_A_scale,
|
|
B_scale=w1_scale,
|
|
use_fp8=use_fp8,
|
|
per_channel_quant=per_channel_quant,
|
|
block_shape=block_shape,
|
|
)
|
|
|
|
- intermediate_cache1_flat = intermediate_cache1.view(-1, intermediate_cache1.shape[-1])
|
|
- intermediate_cache2_input = intermediate_cache1_flat[:, :expert_hidden_dim_per_partition].contiguous()
|
|
+ torch.ops._C.silu_and_mul(
|
|
+ intermediate_cache2,
|
|
+ intermediate_cache1.view(-1, intermediate_cache1.shape[-1]),
|
|
+ )
|
|
+ second_input = intermediate_cache2
|
|
|
|
intermediate_A_scale = None
|
|
if use_fp8:
|
|
@@ -321,13 +334,13 @@ def _run_fused_moe_iteration(
|
|
group_size = block_dims[1] if block_dims else 128
|
|
- intermediate_cache2_input, intermediate_A_scale = quantize_activations_to_fp8(
|
|
- intermediate_cache2_input,
|
|
+ second_input, intermediate_A_scale = quantize_activations_to_fp8(
|
|
+ intermediate_cache2,
|
|
group_size=group_size,
|
|
)
|
|
|
|
_invoke_kernel(
|
|
- A=intermediate_cache2_input,
|
|
+ A=second_input,
|
|
B=w2.contiguous(),
|
|
- C=intermediate_cache2.contiguous(),
|
|
+ C=intermediate_cache3.contiguous(),
|
|
topk_weights=topk_weights.contiguous(),
|
|
sorted_token_ids=sorted_token_ids.contiguous(),
|
|
expert_ids=expert_ids.contiguous(),
|
|
@@ -335,4 +350,6 @@ def _run_fused_moe_iteration(
|
|
)
|
|
+
|
|
+ ops.moe_sum(intermediate_cache3, output)
|
|
|
|
|
|
def _collect_cuda_event_stats(step_fn, active_steps: int) -> Dict:
|
|
@@ -493,6 +508,5 @@ def profile_fused_moe_kernel(
|
|
w1_scale = None
|
|
w2_scale = None
|
|
- A_scale = None
|
|
|
|
block_dims = _validate_block_shape(block_shape)
|
|
if use_fp8:
|
|
@@ -509,10 +521,8 @@ def profile_fused_moe_kernel(
|
|
per_channel=per_channel_quant,
|
|
block_shape=block_shape,
|
|
)
|
|
- group_size = block_dims[1] if block_dims else 128
|
|
- A, A_scale = quantize_activations_to_fp8(A, group_size=group_size)
|
|
|
|
- config_dtype = get_config_dtype_str(base_dtype)
|
|
+ config_dtype = get_config_dtype_str(base_dtype, use_fp8_w8a8=use_fp8)
|
|
config = try_get_optimal_moe_config(
|
|
w1_shape=w1.shape,
|
|
w2_shape=w2.shape,
|
|
@@ -535,13 +544,25 @@ def profile_fused_moe_kernel(
|
|
device=device,
|
|
dtype=output_dtype,
|
|
)
|
|
intermediate_cache2 = torch.empty(
|
|
- num_tokens,
|
|
- top_k,
|
|
- hidden_dim,
|
|
+ num_tokens * top_k,
|
|
+ expert_hidden_dim_per_partition,
|
|
+ device=device,
|
|
+ dtype=output_dtype,
|
|
+ )
|
|
+ intermediate_cache3 = torch.empty(
|
|
+ num_tokens,
|
|
+ top_k,
|
|
+ hidden_dim,
|
|
device=device,
|
|
dtype=output_dtype,
|
|
)
|
|
+ output = torch.empty(
|
|
+ num_tokens,
|
|
+ hidden_dim,
|
|
+ device=device,
|
|
+ dtype=output_dtype,
|
|
+ )
|
|
|
|
def _step() -> None:
|
|
_run_fused_moe_iteration(
|
|
@@ -552,6 +571,8 @@ def profile_fused_moe_kernel(
|
|
w2=w2,
|
|
intermediate_cache1=intermediate_cache1,
|
|
intermediate_cache2=intermediate_cache2,
|
|
+ intermediate_cache3=intermediate_cache3,
|
|
+ output=output,
|
|
topk_weights=topk_weights,
|
|
sorted_token_ids=sorted_token_ids,
|
|
expert_ids=expert_ids,
|
|
@@ -562,6 +583,6 @@ def profile_fused_moe_kernel(
|
|
expert_hidden_dim_per_partition=expert_hidden_dim_per_partition,
|
|
block_dims=block_dims,
|
|
- A_scale=A_scale,
|
|
+ A_scale=None,
|
|
w1_scale=w1_scale,
|
|
w2_scale=w2_scale,
|
|
use_fp8=use_fp8,
|
|
diff --git a/frontier/profiling/moe/moe_impl.py b/frontier/profiling/moe/moe_impl.py
|
|
--- a/frontier/profiling/moe/moe_impl.py
|
|
+++ b/frontier/profiling/moe/moe_impl.py
|
|
@@ -245,10 +245,12 @@ class MoETokenShuffler(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_experts: int,
|
|
router_topk: int,
|
|
hidden_dim: int,
|
|
expert_hidden_dim: int,
|
|
dtype: torch.dtype,
|
|
use_gated: bool,
|
|
num_local_experts: Optional[int] = None,
|
|
+ use_fp8: bool = False,
|
|
+ block_shape: Optional[list[int]] = None,
|
|
):
|
|
@@ -264,6 +266,8 @@ class MoETokenShuffler(nn.Module):
|
|
self.router_topk = router_topk
|
|
self.hidden_dim = hidden_dim
|
|
self.expert_hidden_dim = expert_hidden_dim
|
|
self.dtype = dtype
|
|
self.use_gated = use_gated
|
|
+ self.use_fp8 = use_fp8
|
|
+ self.block_shape = block_shape
|
|
self._block_size_cache = {}
|
|
@@ -325,9 +329,12 @@ class MoETokenShuffler(nn.Module):
|
|
- config_dtype = get_config_dtype_str(dtype=self.dtype)
|
|
+ config_dtype = get_config_dtype_str(
|
|
+ dtype=self.dtype,
|
|
+ use_fp8_w8a8=self.use_fp8,
|
|
+ )
|
|
config = try_get_optimal_moe_config(
|
|
w1_shape=w1_shape,
|
|
w2_shape=w2_shape,
|
|
top_k=self.router_topk,
|
|
dtype=config_dtype,
|
|
M=num_tokens,
|
|
- block_shape=None,
|
|
+ block_shape=self.block_shape,
|
|
)
|
|
diff --git a/frontier/profiling/moe/moe_wrapper.py b/frontier/profiling/moe/moe_wrapper.py
|
|
--- a/frontier/profiling/moe/moe_wrapper.py
|
|
+++ b/frontier/profiling/moe/moe_wrapper.py
|
|
@@ -149,9 +149,11 @@ class MoEWrapper:
|
|
self.shuffler = MoETokenShuffler(
|
|
num_experts=self.num_experts,
|
|
num_local_experts=self.num_experts_per_device,
|
|
router_topk=self.router_topk,
|
|
hidden_dim=self.hidden_dim,
|
|
expert_hidden_dim=self.expert_hidden_dim,
|
|
dtype=self._dtype,
|
|
use_gated=self.use_gated,
|
|
+ use_fp8=self.use_fp8,
|
|
+ block_shape=self.block_shape,
|
|
).to(dtype=self._dtype).cuda().eval()
|