gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Align Frontier FP8 profiling with vLLM runtime

Browse Source
This commit is contained in:
Gahow Wang
2026-07-15 18:34:13 +08:00
parent 9c8570f36b
commit 684a2de413
3 changed files with 525 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
runs/frontier-multicase-sufficiency-v0/best_effort/frontier_fp8_profile_contract.patch Normal file
View File

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

216
runs/frontier-multicase-sufficiency-v0/best_effort/frontier_moe_serving_path.patch Normal file
View File

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

220
runs/frontier-multicase-sufficiency-v0/best_effort/moe_path_parity.py Normal file
View File

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