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()