From 69fdd7a9069ed03451d3586764be207a839850e0 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Tue, 21 Jul 2026 20:38:09 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- .../llm/custom_ops/test_quantized_moe.py | 119 ++++++++++++++++-- 1 file changed, 111 insertions(+), 8 deletions(-) diff --git a/extension/llm/custom_ops/test_quantized_moe.py b/extension/llm/custom_ops/test_quantized_moe.py index 130fe8d331f..0b4299e8f12 100644 --- a/extension/llm/custom_ops/test_quantized_moe.py +++ b/extension/llm/custom_ops/test_quantized_moe.py @@ -90,13 +90,14 @@ def _build_moe_eager( def _moe_forward_with_qdq_weights( moe: MOEFeedForward, x: torch.Tensor, group_size: int ) -> torch.Tensor: - """Run `moe.forward(x)` after replacing each per-expert weight with its - INT4 group-quantize/dequant round-trip. This is the apples-to-apples - reference for the custom op's output. - - We deepcopy the real module (so any attributes added by future - `MOEFeedForward.__init__` / `ConditionalFeedForward.__init__` changes - are preserved) and only overwrite the per-expert weights in-place. + """Independent q-dq reference for `llama::quantized_moe_ffn`. + + Applies the same INT4 group-quantize/dequant round-trip to each expert + weight, then reproduces the op's routing contract (see op_moe.cpp) in + eager PyTorch: sigmoid or softmax scoring, expert-bias-shifted top-k + selection, un-biased weight renormalization scaled by route_scale for the + sigmoid path (softmax path softmaxes the top-k raw scores). Expert compute + reuses the q-dq `ConditionalFeedForward`. """ moe_q = copy.deepcopy(moe) @@ -119,7 +120,109 @@ def _qdq_per_expert(w_eFD: torch.Tensor) -> torch.Tensor: w2_DF_qdq = _qdq_per_expert(w2_DF) cond_q.w2.copy_(w2_DF_qdq.transpose(-2, -1).contiguous()) - return moe_q.forward(x) + scores = moe.gate(x) # [T, E] + k = moe.num_activated_experts + if moe.score_func == "sigmoid": + s = torch.sigmoid(scores) + sel = s + moe.expert_bias if getattr(moe, "use_expert_bias", False) else s + idx = torch.topk(sel, k, dim=-1).indices + weights = torch.gather(s, -1, idx) + weights = ( + weights * moe.route_scale / (weights.sum(dim=-1, keepdim=True) + 1e-20) + ) + else: + idx = torch.topk(scores, k, dim=-1).indices + weights = torch.gather(scores, -1, idx).softmax(dim=-1) + + expert_outs = moe_q.cond_ffn(x, idx) # [T, K, D] + return torch.einsum("tkd,tk->td", expert_outs, weights) + + +class TestQuantizedMoeFfnOp(unittest.TestCase): + """Numerical correctness vs a Python q-dq reference.""" + + def setUp(self) -> None: + torch.manual_seed(0) + + def _check_against_qdq_reference( + self, + *, + score_func: str, + use_expert_bias: bool, + route_scale: float, + num_tokens: int = 8, + atol: float = 5e-3, + ) -> None: + dim, hidden_dim = 32, 32 + num_experts, num_activated_experts = 4, 2 + group_size = 32 + + moe = _build_moe_eager( + dim=dim, + hidden_dim=hidden_dim, + num_experts=num_experts, + num_activated_experts=num_activated_experts, + score_func=score_func, + use_expert_bias=use_expert_bias, + route_scale=route_scale, + ) + x = torch.randn(num_tokens, dim) + + with torch.no_grad(): + ref = _moe_forward_with_qdq_weights(moe, x, group_size) + + qmodel = MOEFeedForward.__new__(MOEFeedForward) + torch.nn.Module.__init__(qmodel) + qmodel.gate = moe.gate + qmodel.cond_ffn = moe.cond_ffn + qmodel.dim = moe.dim + qmodel.num_activated_experts = moe.num_activated_experts + qmodel.score_func = moe.score_func + qmodel.route_scale = moe.route_scale + qmodel.use_expert_bias = moe.use_expert_bias + if moe.use_expert_bias: + qmodel.expert_bias = moe.expert_bias + + wrapper = torch.nn.Module() + wrapper.block_sparse_moe = qmodel + replace_moe_with_quantized_op(wrapper, group_size=group_size, weight_nbit=4) + test = wrapper.block_sparse_moe(x) + + diff = (ref - test).abs() + self.assertTrue( + diff.max().item() < atol, + f"max abs diff {diff.max().item()} > atol {atol} (mean={diff.mean().item()})", + ) + + def test_sigmoid_with_bias_route_scale_2p5(self) -> None: + self._check_against_qdq_reference( + score_func="sigmoid", use_expert_bias=True, route_scale=2.5 + ) + + def test_sigmoid_no_bias(self) -> None: + self._check_against_qdq_reference( + score_func="sigmoid", use_expert_bias=False, route_scale=1.0 + ) + + def test_softmax(self) -> None: + self._check_against_qdq_reference( + score_func="softmax", use_expert_bias=False, route_scale=1.0 + ) + + def test_single_token(self) -> None: + self._check_against_qdq_reference( + score_func="sigmoid", + use_expert_bias=True, + route_scale=2.5, + num_tokens=1, + ) + + def test_route_scale_zero(self) -> None: + self._check_against_qdq_reference( + score_func="sigmoid", + use_expert_bias=False, + route_scale=0.0, + ) class TestSourceTransform(unittest.TestCase): From aa32a3499f64b200743c93458e04e3ee703d2e0f Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Thu, 23 Jul 2026 13:58:48 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- extension/llm/custom_ops/test_quantized_moe.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/extension/llm/custom_ops/test_quantized_moe.py b/extension/llm/custom_ops/test_quantized_moe.py index 6db5065df1a..32706aa52f9 100644 --- a/extension/llm/custom_ops/test_quantized_moe.py +++ b/extension/llm/custom_ops/test_quantized_moe.py @@ -105,18 +105,10 @@ def _make_valid_meta_inputs() -> dict: "gate_weight": torch.empty( (num_experts, dim), dtype=torch.float32, device="meta" ), - "expert_bias": torch.empty( - (num_experts,), dtype=torch.float32, device="meta" - ), - "packed_w1": torch.empty( - (num_experts, 1), dtype=torch.uint8, device="meta" - ), - "packed_w3": torch.empty( - (num_experts, 1), dtype=torch.uint8, device="meta" - ), - "packed_w2": torch.empty( - (num_experts, 1), dtype=torch.uint8, device="meta" - ), + "expert_bias": torch.empty((num_experts,), dtype=torch.float32, device="meta"), + "packed_w1": torch.empty((num_experts, 1), dtype=torch.uint8, device="meta"), + "packed_w3": torch.empty((num_experts, 1), dtype=torch.uint8, device="meta"), + "packed_w2": torch.empty((num_experts, 1), dtype=torch.uint8, device="meta"), "num_activated_experts": 2, "num_experts": num_experts, "hidden_dim": hidden_dim, @@ -179,6 +171,7 @@ def _qdq_per_expert(w_eFD: torch.Tensor) -> torch.Tensor: return torch.einsum("tkd,tk->td", expert_outs, weights) +@_REQUIRES_TORCHAO_KERNEL_LIBRARY class TestQuantizedMoeFfnOp(unittest.TestCase): """Numerical correctness vs a Python q-dq reference."""