diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_fp8_profile_contract.patch b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_fp8_profile_contract.patch new file mode 100644 index 0000000..012fd79 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_fp8_profile_contract.patch @@ -0,0 +1,89 @@ +diff --git a/frontier/config/quantization_manager.py b/frontier/config/quantization_manager.py +--- a/frontier/config/quantization_manager.py ++++ b/frontier/config/quantization_manager.py +@@ -356,9 +356,17 @@ class QuantizationManager: + target_precision = self.get_precision(op_name) + precision_match = target_precision == profiling_precision + quant_match = profiling_quant_signature == expected_quant_signature ++ # A mixed-precision profiling CSV records the model's output dtype at ++ # file level. Quantized operators in that CSV were nevertheless ++ # measured with the exact quantization scheme identified by the ++ # quant signature. Treating those samples as BF16 and scaling them ++ # again would double-apply the FP8 speedup. ++ exact_fp8_profile = quant_match and target_precision == PrecisionType.FP8 + with self._lock: +- self._operation_profiling_precision[op_name] = profiling_precision +- if precision_match: ++ self._operation_profiling_precision[op_name] = ( ++ target_precision if exact_fp8_profile else profiling_precision ++ ) ++ if precision_match or exact_fp8_profile: + self._operation_data_sources[op_name] = "profiling" + self._operation_approximation_factors.pop(op_name, None) + self._operation_speedup_factors.pop(op_name, None) +diff --git a/tests/unit/test_quantization_profile_contract.py b/tests/unit/test_quantization_profile_contract.py +new file mode 100644 +--- /dev/null ++++ b/tests/unit/test_quantization_profile_contract.py +@@ -0,0 +1,61 @@ ++from unittest.mock import MagicMock ++ ++import pytest ++ ++from frontier.config.model_config import QuantizationConfig ++from frontier.config.precision_type import PrecisionType ++from frontier.config.quantization_manager import QuantizationManager ++ ++ ++def test_exact_fp8_quant_signature_does_not_rescale_mixed_profile() -> None: ++ manager = QuantizationManager() ++ manager.load_config() ++ ++ quant_config = QuantizationConfig( ++ quant_method="fp8", ++ activation_scheme="dynamic", ++ is_checkpoint_fp8_serialized=True, ++ weight_block_size=(128, 128), ++ ) ++ quant_signature = quant_config.get_quant_signature() ++ model_config = MagicMock() ++ model_config.get_default_precision.return_value = PrecisionType.BF16 ++ model_config.get_name.return_value = "Qwen3-235B-A22B-FP8" ++ model_config.torch_dtype = "bfloat16" ++ model_config.quantization_config = quant_config ++ model_config.get_quant_signature.return_value = quant_signature ++ ++ manager.configure_from_model_config(model_config) ++ manager.register_profiling_metadata( ++ operation_names=["attn_pre_proj"], ++ profiling_precision=PrecisionType.BF16, ++ profiling_quant_signature=quant_signature, ++ expected_quant_signature=quant_signature, ++ file_path="linear_op.csv", ++ ) ++ ++ metadata = { ++ item["operation"]: item ++ for item in manager.get_operation_precision_metadata() ++ } ++ assert metadata["attn_pre_proj"]["data_source"] == "profiling" ++ assert metadata["attn_pre_proj"]["approximation_factor"] is None ++ assert manager.has_precision_mismatch("attn_pre_proj") is False ++ assert manager.adjust_compute_time("attn_pre_proj", 1.25) == pytest.approx(1.25) ++ ++ ++def test_mismatched_fp8_quant_signature_still_uses_approximation() -> None: ++ manager = QuantizationManager() ++ manager.load_config() ++ manager._operation_precisions["attn_pre_proj"] = PrecisionType.FP8 ++ ++ manager.register_profiling_metadata( ++ operation_names=["attn_pre_proj"], ++ profiling_precision=PrecisionType.BF16, ++ profiling_quant_signature="none", ++ expected_quant_signature="method=fp8", ++ file_path="linear_op.csv", ++ ) ++ ++ assert manager.has_precision_mismatch("attn_pre_proj") is True ++ assert manager.adjust_compute_time("attn_pre_proj", 1.0) == pytest.approx(0.5) diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_path.patch b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_path.patch new file mode 100644 index 0000000..6702a90 --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_path.patch @@ -0,0 +1,216 @@ +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() diff --git a/runs/frontier-multicase-sufficiency-v0/best_effort/moe_path_parity.py b/runs/frontier-multicase-sufficiency-v0/best_effort/moe_path_parity.py new file mode 100644 index 0000000..e90654f --- /dev/null +++ b/runs/frontier-multicase-sufficiency-v0/best_effort/moe_path_parity.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +"""Compare patched Frontier MoE decomposition with vLLM's serving path.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from frontier.profiling.moe.moe_vllm_kernel import ( + profile_fused_moe_kernel, + quantize_weights_to_fp8, +) +from vllm.model_executor.layers.fused_moe.fused_moe import ( + fused_experts, + get_config_dtype_str, + moe_align_block_size, + try_get_optimal_moe_config, +) + + +def _measure(step, warmup: int, active: int) -> dict[str, float]: + for _ in range(warmup): + step() + torch.cuda.synchronize() + + samples = [] + for _ in range(active): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + step() + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + values = torch.tensor(samples) + return { + "min": float(values.min()), + "median": float(values.median()), + "mean": float(values.mean()), + "max": float(values.max()), + "std": float(values.std()), + } + + +def _routing(num_tokens: int, top_k: int, num_experts: int, seed: int): + generator = torch.Generator(device="cuda") + generator.manual_seed(seed) + topk_ids = torch.randint( + num_experts, + (num_tokens, top_k), + generator=generator, + device="cuda", + dtype=torch.int64, + ) + topk_weights = torch.rand( + (num_tokens, top_k), + generator=generator, + device="cuda", + dtype=torch.float32, + ) + topk_weights /= topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--tokens", nargs="+", type=int, default=[16, 256, 1024]) + parser.add_argument("--tp", type=int, default=4) + parser.add_argument("--ep", type=int, default=1) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--active", type=int, default=20) + parser.add_argument("--seed", type=int, default=20260715) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + hidden_dim = 4096 + expert_hidden_dim = 1536 + global_num_experts = 128 + top_k = 8 + block_shape = [128, 128] + if global_num_experts % args.ep: + raise ValueError("EP must divide 128 experts") + if expert_hidden_dim % args.tp: + raise ValueError("TP must divide the expert intermediate dimension") + + local_num_experts = global_num_experts // args.ep + local_intermediate = expert_hidden_dim // args.tp + device = torch.device("cuda") + + torch.manual_seed(args.seed) + w1_bf16 = torch.randn( + local_num_experts, + 2 * local_intermediate, + hidden_dim, + dtype=torch.bfloat16, + device=device, + ) + w2_bf16 = torch.randn( + local_num_experts, + hidden_dim, + local_intermediate, + dtype=torch.bfloat16, + device=device, + ) + w1, w1_scale = quantize_weights_to_fp8(w1_bf16, block_shape=block_shape) + w2, w2_scale = quantize_weights_to_fp8(w2_bf16, block_shape=block_shape) + del w1_bf16, w2_bf16 + torch.cuda.empty_cache() + + rows = [] + for index, num_tokens in enumerate(args.tokens): + topk_weights, topk_ids = _routing( + num_tokens, + top_k, + global_num_experts, + args.seed + index, + ) + hidden_states = torch.randn( + num_tokens, + hidden_dim, + dtype=torch.bfloat16, + device=device, + ) + + frontier_grouped = profile_fused_moe_kernel( + num_tokens=num_tokens, + num_experts=local_num_experts, + hidden_dim=hidden_dim, + expert_hidden_dim=expert_hidden_dim, + top_k=top_k, + topk_weights=topk_weights, + topk_ids=topk_ids, + tensor_parallel_size=args.tp, + dtype=torch.bfloat16, + warmup_steps=args.warmup, + active_steps=args.active, + use_fp8=True, + per_channel_quant=False, + block_shape=block_shape, + global_num_experts=global_num_experts, + ) + + config = try_get_optimal_moe_config( + w1_shape=w1.shape, + w2_shape=w2.shape, + top_k=top_k, + dtype=get_config_dtype_str( + torch.bfloat16, + use_fp8_w8a8=True, + ), + M=num_tokens, + block_shape=block_shape, + ) + + def align_step() -> None: + moe_align_block_size( + topk_ids, + config["BLOCK_SIZE_M"], + global_num_experts, + ) + + alignment = _measure(align_step, args.warmup, args.active) + + def serving_step() -> None: + fused_experts( + hidden_states=hidden_states, + w1=w1, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + inplace=False, + use_fp8_w8a8=True, + per_channel_quant=False, + global_num_experts=global_num_experts, + w1_scale=w1_scale, + w2_scale=w2_scale, + block_shape=block_shape, + ) + + serving = _measure(serving_step, args.warmup, args.active) + decomposed_ms = frontier_grouped["median"] + alignment["median"] + rows.append( + { + "num_tokens": num_tokens, + "block_size_m": config["BLOCK_SIZE_M"], + "frontier_grouped_ms": frontier_grouped, + "frontier_alignment_ms": alignment, + "frontier_decomposed_median_ms": decomposed_ms, + "vllm_fused_experts_ms": serving, + "decomposed_over_serving": decomposed_ms / serving["median"], + } + ) + + result = { + "contract": "Frontier grouped_gemm + shuffling alignment vs vLLM fused_experts", + "model_shape": { + "hidden_dim": hidden_dim, + "expert_hidden_dim": expert_hidden_dim, + "global_num_experts": global_num_experts, + "local_num_experts": local_num_experts, + "top_k": top_k, + "tp": args.tp, + "ep": args.ep, + "dtype": "block_fp8_w8a8_bf16_output", + "block_shape": block_shape, + }, + "warmup": args.warmup, + "active": args.active, + "rows": rows, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main()