From f8851e0a53477021590fb69666e1606fc4fd7c09 Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 19:01:47 -0500
Subject: [PATCH 01/14] stash

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../layers/fused_moe/prepare_finalize.py      | 86 ++++++++++++++++++-
 1 file changed, 85 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index 5d806fa843a3..5b9a7d65c366 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -11,12 +11,96 @@
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 
+class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize):
+    def __init__(self) -> None:
+        super().__init__()
+        self.dummy_tensor = torch.empty(1, device='cuda')
+
+    @property
+    def activation_format(self) -> mk.FusedMoEActivationFormat:
+        return mk.FusedMoEActivationFormat.Standard
+
+    def max_num_tokens_per_rank(self) -> int | None:
+        return None
+
+    def topk_indices_dtype(self) -> torch.dtype | None:
+        return None
+
+    def num_dispatchers(self) -> int:
+        return 1
+
+    def output_is_reduced(self) -> bool:
+        return False
+    
+    def prepare(
+        self,
+        a1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        num_experts: int,
+        expert_map: torch.Tensor | None,
+        apply_router_weight_on_input: bool,
+        quant_config: FusedMoEQuantConfig,
+    ) -> mk.PrepareResultType:
+        from vllm.distributed import get_ep_group
+
+        if not hasattr(self, "dummy_tensor"):
+            self.dummy_tensor = torch.zeros(1, device='cuda')
+
+        extra_tensors = [topk_weights,topk_ids]
+        a1, _, extra_tensors = get_ep_group().dispatch(
+            a1,
+            self.dummy_tensor,
+            is_sequence_parallel=False, # TODO?
+            extra_tensors=extra_tensors,
+        )
+        topk_weights, topk_ids = extra_tensors
+
+        a1q, a1q_scale = moe_kernel_quantize_input(
+            a1,
+            quant_config.a1_scale,
+            quant_config.quant_dtype,
+            quant_config.per_act_token_quant,
+            quant_config.block_shape,
+        )
+        expert_tokens_meta = None # TODO?
+
+        return a1q, a1q_scale, None, topk_weights, topk_ids
+
+    def finalize(
+        self,
+        output: torch.Tensor,
+        fused_expert_output: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        apply_router_weight_on_input: bool,
+        weight_and_reduce_impl: mk.TopKWeightAndReduce,
+    ) -> None:
+        if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
+            weight_and_reduce_impl = TopKWeightAndReduceContiguous()
+        weight_and_reduce_impl.apply(
+            output=output,
+            fused_expert_output=fused_expert_output,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+        )
+
+        from vllm.distributed import get_ep_group
+        combined_output = get_ep_group().combine(
+            output,
+            is_sequence_parallel=False
+        )
+
+        combined_output.copy_(output)
+
+
 
 class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
     def __init__(self, defer_input_quant: bool = False) -> None:
         super().__init__()
         self.defer_input_quant = defer_input_quant
-
+    
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard

From 085adf7377692bfe287dd58c8e5d98b2dad32e18 Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 19:02:06 -0500
Subject: [PATCH 02/14] stash

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/model_executor/layers/fused_moe/prepare_finalize.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index 5b9a7d65c366..06bcb4f68b22 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -50,7 +50,7 @@ def prepare(
         extra_tensors = [topk_weights,topk_ids]
         a1, _, extra_tensors = get_ep_group().dispatch(
             a1,
-            self.dummy_tensor,
+            self.dummy_tensor, # router logits
             is_sequence_parallel=False, # TODO?
             extra_tensors=extra_tensors,
         )
@@ -91,7 +91,6 @@ def finalize(
             output,
             is_sequence_parallel=False
         )
-
         combined_output.copy_(output)
 
 

From a6b039dd691bfd74fb017a807cb0e39bf0270a7f Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 19:14:29 -0500
Subject: [PATCH 03/14] update interface

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../device_communicators/all2all.py           | 44 ++++++++++++-------
 .../base_device_communicator.py               | 10 +++--
 .../device_communicators/cuda_communicator.py | 10 +++--
 .../device_communicators/xpu_communicator.py  |  8 ++--
 .../model_executor/layers/quantization/fp8.py |  4 ++
 5 files changed, 49 insertions(+), 27 deletions(-)

diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py
index 7a4e81cf967d..891c9439da25 100644
--- a/vllm/distributed/device_communicators/all2all.py
+++ b/vllm/distributed/device_communicators/all2all.py
@@ -62,10 +62,11 @@ def naive_multicast(
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if extra_tensors is not None:
             raise NotImplementedError(
                 "extra_tensors is not supported for NaiveAll2AllManager"
@@ -78,11 +79,14 @@ def dispatch(
         hidden_states = self.naive_multicast(
             hidden_states, cu_tokens_across_sp_cpu, is_sequence_parallel
         )
-        router_logits = self.naive_multicast(
-            router_logits, cu_tokens_across_sp_cpu, is_sequence_parallel
+        topk_weights = self.naive_multicast(
+            topk_weights, cu_tokens_across_sp_cpu, is_sequence_parallel
+        )
+        topk_ids = self.naive_multicast(
+            topk_ids, cu_tokens_across_sp_cpu, is_sequence_parallel
         )
 
-        return hidden_states, router_logits
+        return hidden_states, topk_weights, topk_ids
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
@@ -117,12 +121,13 @@ def __init__(self, cpu_group):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> (
-        tuple[torch.Tensor, torch.Tensor]
-        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         """
         Gather hidden_states and router_logits from all dp ranks.
@@ -134,7 +139,7 @@ def dispatch(
         dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
 
-        tensors_to_gather = [hidden_states, router_logits]
+        tensors_to_gather = [hidden_states, topk_weights, topk_ids]
         if extra_tensors is not None:
             tensors_to_gather.extend(extra_tensors)
 
@@ -144,9 +149,14 @@ def dispatch(
             sizes=sizes,
         )
 
-        if extra_tensors is not None:
-            return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
-        return gathered_tensors[0], gathered_tensors[1]
+        hidden_states = gathered_tensors[0]
+        topk_weights = gathered_tensors[1]
+        topk_ids = gathered_tensors[2]
+
+        if extra_tensors is None:
+            return hidden_states, topk_weights, topk_ids
+
+        return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
@@ -219,10 +229,11 @@ def get_handle(self, kwargs):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         raise NotImplementedError
 
     def combine(
@@ -267,10 +278,11 @@ def get_handle(self, kwargs):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         raise NotImplementedError
 
     def combine(
diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py
index 4a2a7ec5b728..3a951e4be23b 100644
--- a/vllm/distributed/device_communicators/base_device_communicator.py
+++ b/vllm/distributed/device_communicators/base_device_communicator.py
@@ -67,7 +67,8 @@ def get_handle(self, kwargs):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> Any:
@@ -283,15 +284,16 @@ def prepare_communication_buffer_for_model(self, model: torch.nn.Module) -> None
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Dispatch the hidden states and router logits to the appropriate device.
         This is a no-op in the base class.
         """
-        return hidden_states, router_logits
+        return hidden_states, topk_weights, topk_ids
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
index 9542498c453e..6d725be241e1 100644
--- a/vllm/distributed/device_communicators/cuda_communicator.py
+++ b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -321,17 +321,19 @@ def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None):
     def dispatch(  # type: ignore[override]
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> (
-        tuple[torch.Tensor, torch.Tensor]
-        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch(
             hidden_states,
-            router_logits,
+            topk_weights,
+            topk_ids,
             is_sequence_parallel,
             extra_tensors,  # type: ignore[call-arg]
         )
diff --git a/vllm/distributed/device_communicators/xpu_communicator.py b/vllm/distributed/device_communicators/xpu_communicator.py
index f3d9262d20cf..d7a9d832370f 100644
--- a/vllm/distributed/device_communicators/xpu_communicator.py
+++ b/vllm/distributed/device_communicators/xpu_communicator.py
@@ -76,14 +76,16 @@ def broadcast(self, input_: torch.Tensor, src: int = 0) -> None:
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         assert self.all2all_manager is not None
         return self.all2all_manager.dispatch(
             hidden_states,
-            router_logits,
+            topk_weights,
+            topk_ids,
             is_sequence_parallel,
             extra_tensors,  # type: ignore[call-arg]
         )
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index 1223c6902e5f..003a683e26aa 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -1323,10 +1323,14 @@ def apply(
                     apply_router_weight_on_input=layer.apply_router_weight_on_input,
                 )
 
+        logger.info_once(f"{router_logits.shape=}")
         topk_weights, topk_ids = layer.select_experts(
             hidden_states=x,
             router_logits=router_logits,
         )
+        logger.info_once(f"{router_logits.shape=}")
+        logger.info_once(f"{topk_weights.shape=}")
+
         result = self.kernel(
             x,
             layer.w13_weight,

From f8052ce5b8e9d7072776b78397542b4c2c286d1f Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 19:37:47 -0500
Subject: [PATCH 04/14] stash

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/distributed/parallel_state.py            | 12 +++---
 .../layers/fused_moe/prepare_finalize.py      | 39 +++++++++++--------
 .../model_executor/layers/quantization/fp8.py |  5 +--
 3 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py
index 4611d42a5874..f5a9aea4a285 100644
--- a/vllm/distributed/parallel_state.py
+++ b/vllm/distributed/parallel_state.py
@@ -1005,22 +1005,24 @@ def prepare_communication_buffer_for_model(self, model: torch.nn.Module):
     def dispatch(
         self,
         hidden_states: torch.Tensor,
-        router_logits: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
         is_sequence_parallel: bool = False,
         extra_tensors: list[torch.Tensor] | None = None,
     ) -> (
-        tuple[torch.Tensor, torch.Tensor]
-        | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         if self.device_communicator is not None:
             return self.device_communicator.dispatch(  # type: ignore[call-arg]
                 hidden_states,
-                router_logits,
+                topk_weights,
+                topk_ids,
                 is_sequence_parallel,
                 extra_tensors,
             )
         else:
-            return hidden_states, router_logits
+            return hidden_states, topk_weights, topk_ids
 
     def combine(
         self, hidden_states, is_sequence_parallel: bool = False
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index 06bcb4f68b22..1f036b5e227c 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -11,10 +11,11 @@
 )
 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
 
+
 class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize):
     def __init__(self) -> None:
         super().__init__()
-        self.dummy_tensor = torch.empty(1, device='cuda')
+        self.dummy_tensor = torch.empty(1, device="cuda")
 
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -31,7 +32,7 @@ def num_dispatchers(self) -> int:
 
     def output_is_reduced(self) -> bool:
         return False
-    
+
     def prepare(
         self,
         a1: torch.Tensor,
@@ -44,17 +45,20 @@ def prepare(
     ) -> mk.PrepareResultType:
         from vllm.distributed import get_ep_group
 
-        if not hasattr(self, "dummy_tensor"):
-            self.dummy_tensor = torch.zeros(1, device='cuda')
+        print(f"before: {a1.shape=}")
+        print(f"before: {topk_weights.shape=}")
+        print(f"before: {topk_ids.shape=}")
 
-        extra_tensors = [topk_weights,topk_ids]
-        a1, _, extra_tensors = get_ep_group().dispatch(
+        a1, topk_weights, topk_ids = get_ep_group().dispatch(
             a1,
-            self.dummy_tensor, # router logits
-            is_sequence_parallel=False, # TODO?
-            extra_tensors=extra_tensors,
+            topk_weights,
+            topk_ids,
+            # TODO?
+            is_sequence_parallel=False,
         )
-        topk_weights, topk_ids = extra_tensors
+        print(f"after: {a1.shape=}")
+        print(f"after: {topk_weights.shape=}")
+        print(f"after: {topk_ids.shape=}")
 
         a1q, a1q_scale = moe_kernel_quantize_input(
             a1,
@@ -63,7 +67,11 @@ def prepare(
             quant_config.per_act_token_quant,
             quant_config.block_shape,
         )
-        expert_tokens_meta = None # TODO?
+        # TODO?
+        expert_tokens_meta = None
+        print(
+            f"{a1q.dtype=}, {a1q_scale.dtype=} {topk_weights.dtype=}, {topk_ids.dtype=}"
+        )
 
         return a1q, a1q_scale, None, topk_weights, topk_ids
 
@@ -87,19 +95,16 @@ def finalize(
         )
 
         from vllm.distributed import get_ep_group
-        combined_output = get_ep_group().combine(
-            output,
-            is_sequence_parallel=False
-        )
-        combined_output.copy_(output)
 
+        combined_output = get_ep_group().combine(output, is_sequence_parallel=False)
+        combined_output.copy_(output)
 
 
 class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
     def __init__(self, defer_input_quant: bool = False) -> None:
         super().__init__()
         self.defer_input_quant = defer_input_quant
-    
+
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index 003a683e26aa..610ebe3b3399 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -1123,6 +1123,7 @@ def maybe_make_prepare_finalize(
         self,
         routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     ) -> mk.FusedMoEPrepareAndFinalize | None:
+        # return MoEPrepareAndFinalizeNaiveEP()
         if (
             self.fp8_backend == Fp8MoeBackend.AITER
             or self.fp8_backend == Fp8MoeBackend.MARLIN
@@ -1323,13 +1324,11 @@ def apply(
                     apply_router_weight_on_input=layer.apply_router_weight_on_input,
                 )
 
-        logger.info_once(f"{router_logits.shape=}")
         topk_weights, topk_ids = layer.select_experts(
             hidden_states=x,
             router_logits=router_logits,
         )
-        logger.info_once(f"{router_logits.shape=}")
-        logger.info_once(f"{topk_weights.shape=}")
+        print(f"{x.dtype=}, {topk_weights.dtype=}, {topk_ids.dtype=}")
 
         result = self.kernel(
             x,

From 13b619ff7d04d6e579c39bfbcf10cd61153f2d1d Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 20:20:23 -0500
Subject: [PATCH 05/14] stash

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../layers/fused_moe/fused_moe_modular_method.py            | 1 +
 vllm/model_executor/layers/quantization/fp8.py              | 6 +++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
index 6abefde0763e..53a132351374 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
@@ -95,6 +95,7 @@ def apply(
             hidden_states=x,
             router_logits=router_logits,
         )
+        logger.info(f"after se: {x.dtype=}, {topk_weights.dtype=}, {topk_ids.dtype=}")
 
         result = self.fused_experts(
             hidden_states=x,
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index 610ebe3b3399..d612ada0449d 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -1123,7 +1123,11 @@ def maybe_make_prepare_finalize(
         self,
         routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     ) -> mk.FusedMoEPrepareAndFinalize | None:
-        # return MoEPrepareAndFinalizeNaiveEP()
+        from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+            MoEPrepareAndFinalizeNaiveEP,
+        )
+
+        return MoEPrepareAndFinalizeNaiveEP()
         if (
             self.fp8_backend == Fp8MoeBackend.AITER
             or self.fp8_backend == Fp8MoeBackend.MARLIN

From 04bb01013b7b19d6b98a2f608e36cceaaf2afce2 Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 20:33:47 -0500
Subject: [PATCH 06/14] first correctness!

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../layers/fused_moe/fused_moe.py             |  3 ++
 .../layers/fused_moe/prepare_finalize.py      | 28 ++++++-------------
 2 files changed, 11 insertions(+), 20 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index c4047401c0e7..7ca65bc85628 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -2315,6 +2315,9 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
+        logger.info(
+            f"in apply: {hidden_states.dtype=}, {topk_weights.dtype=}, {topk_ids.dtype=}, {a1q_scale.dtype=}"
+        )
         # Check constraints.
         if self.quant_config.use_int4_w4a16:
             assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index 1f036b5e227c..de86469d491a 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -4,6 +4,7 @@
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.distributed import get_ep_group
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
     TopKWeightAndReduceContiguous,
@@ -43,12 +44,6 @@ def prepare(
         apply_router_weight_on_input: bool,
         quant_config: FusedMoEQuantConfig,
     ) -> mk.PrepareResultType:
-        from vllm.distributed import get_ep_group
-
-        print(f"before: {a1.shape=}")
-        print(f"before: {topk_weights.shape=}")
-        print(f"before: {topk_ids.shape=}")
-
         a1, topk_weights, topk_ids = get_ep_group().dispatch(
             a1,
             topk_weights,
@@ -56,9 +51,6 @@ def prepare(
             # TODO?
             is_sequence_parallel=False,
         )
-        print(f"after: {a1.shape=}")
-        print(f"after: {topk_weights.shape=}")
-        print(f"after: {topk_ids.shape=}")
 
         a1q, a1q_scale = moe_kernel_quantize_input(
             a1,
@@ -67,13 +59,10 @@ def prepare(
             quant_config.per_act_token_quant,
             quant_config.block_shape,
         )
-        # TODO?
+        # TODO - this is just for deepgemm
         expert_tokens_meta = None
-        print(
-            f"{a1q.dtype=}, {a1q_scale.dtype=} {topk_weights.dtype=}, {topk_ids.dtype=}"
-        )
 
-        return a1q, a1q_scale, None, topk_weights, topk_ids
+        return a1q, a1q_scale, None, topk_ids, topk_weights
 
     def finalize(
         self,
@@ -86,18 +75,17 @@ def finalize(
     ) -> None:
         if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
             weight_and_reduce_impl = TopKWeightAndReduceContiguous()
-        weight_and_reduce_impl.apply(
-            output=output,
+
+        out = weight_and_reduce_impl.apply(
+            output=None,
             fused_expert_output=fused_expert_output,
             topk_weights=topk_weights,
             topk_ids=topk_ids,
             apply_router_weight_on_input=apply_router_weight_on_input,
         )
 
-        from vllm.distributed import get_ep_group
-
-        combined_output = get_ep_group().combine(output, is_sequence_parallel=False)
-        combined_output.copy_(output)
+        # ... copy the
+        output.copy_(get_ep_group().combine(out, is_sequence_parallel=False))
 
 
 class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

From b1320de56e5d06dfe32148bd86bceb847e634f5d Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 20:40:55 -0500
Subject: [PATCH 07/14] updated

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/model_executor/layers/quantization/fp8.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index d612ada0449d..26c371a8a56f 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -1332,7 +1332,6 @@ def apply(
             hidden_states=x,
             router_logits=router_logits,
         )
-        print(f"{x.dtype=}, {topk_weights.dtype=}, {topk_ids.dtype=}")
 
         result = self.kernel(
             x,

From 4d472065615a080e72b53a1265309e1b531c3cef Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 20:41:45 -0500
Subject: [PATCH 08/14] comments

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/model_executor/layers/fused_moe/fused_moe.py              | 3 ---
 .../layers/fused_moe/fused_moe_modular_method.py               | 1 -
 2 files changed, 4 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index 7ca65bc85628..c4047401c0e7 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -2315,9 +2315,6 @@ def apply(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        logger.info(
-            f"in apply: {hidden_states.dtype=}, {topk_weights.dtype=}, {topk_ids.dtype=}, {a1q_scale.dtype=}"
-        )
         # Check constraints.
         if self.quant_config.use_int4_w4a16:
             assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
index 53a132351374..6abefde0763e 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
@@ -95,7 +95,6 @@ def apply(
             hidden_states=x,
             router_logits=router_logits,
         )
-        logger.info(f"after se: {x.dtype=}, {topk_weights.dtype=}, {topk_ids.dtype=}")
 
         result = self.fused_experts(
             hidden_states=x,

From f86fad8c14f1458d08da3890c2e292c7ef6923ad Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 20:42:47 -0500
Subject: [PATCH 09/14] updated

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/model_executor/layers/fused_moe/prepare_finalize.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index de86469d491a..f406df814eff 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -14,10 +14,6 @@
 
 
 class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize):
-    def __init__(self) -> None:
-        super().__init__()
-        self.dummy_tensor = torch.empty(1, device="cuda")
-
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
@@ -59,10 +55,10 @@ def prepare(
             quant_config.per_act_token_quant,
             quant_config.block_shape,
         )
-        # TODO - this is just for deepgemm
+        # TODO - this is just for deepgemm?
         expert_tokens_meta = None
 
-        return a1q, a1q_scale, None, topk_ids, topk_weights
+        return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights
 
     def finalize(
         self,

From 8c1a530d976776aadb69a314c32fbd6f6c8d1f2b Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 21:17:26 -0500
Subject: [PATCH 10/14] updateds

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/model_executor/layers/fused_moe/all2all_utils.py |  7 +++++++
 vllm/model_executor/layers/fused_moe/config.py        | 10 ++++++++++
 vllm/model_executor/layers/quantization/fp8.py        |  6 +-----
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py
index 036b3cac4cb3..d76c629b72ba 100644
--- a/vllm/model_executor/layers/fused_moe/all2all_utils.py
+++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py
@@ -15,6 +15,9 @@
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
     FusedMoEPrepareAndFinalize,
 )
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNaiveEP,
+)
 from vllm.platforms import current_platform
 from vllm.utils.import_utils import has_deep_ep, has_pplx
 
@@ -170,4 +173,8 @@ def maybe_make_prepare_finalize(
             local_expert_global_ids=local_expert_global_ids,
         )
 
+    elif moe.use_naive_kernels:
+        prepare_finalize = MoEPrepareAndFinalizeNaiveEP()
+
+    print(f"{prepare_finalize=}")
     return prepare_finalize
diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py
index 17d5ec4bcda7..162fe36ba044 100644
--- a/vllm/model_executor/layers/fused_moe/config.py
+++ b/vllm/model_executor/layers/fused_moe/config.py
@@ -855,6 +855,12 @@ def use_deepep_ht_kernels(self):
     def use_deepep_ll_kernels(self):
         return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
 
+    @property
+    def use_naive_kernels(self):
+        return self.use_all2all_kernels and (
+            self.all2all_backend in ["allgather_reducescatter", "naive"]
+        )
+
     @staticmethod
     def flatten_tp_across_dp_and_pcp(
         tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
@@ -1076,6 +1082,10 @@ def use_deepep_ht_kernels(self):
     def use_deepep_ll_kernels(self):
         return self.moe_parallel_config.use_deepep_ll_kernels
 
+    @property
+    def use_naive_kernels(self):
+        return self.moe_parallel_config.use_naive_kernels
+
     @property
     def use_flashinfer_cutlass_kernels(self):
         """
diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py
index 80db65c2f774..df494c05ec2d 100644
--- a/vllm/model_executor/layers/quantization/fp8.py
+++ b/vllm/model_executor/layers/quantization/fp8.py
@@ -877,11 +877,6 @@ def maybe_make_prepare_finalize(
         self,
         routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
     ) -> mk.FusedMoEPrepareAndFinalize | None:
-        from vllm.model_executor.layers.fused_moe.prepare_finalize import (
-            MoEPrepareAndFinalizeNaiveEP,
-        )
-
-        return MoEPrepareAndFinalizeNaiveEP()
         if self.fp8_backend in [
             Fp8MoeBackend.AITER,
             Fp8MoeBackend.MARLIN,
@@ -889,6 +884,7 @@ def maybe_make_prepare_finalize(
         ]:
             return None
         elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
+            # TODO(rob): we can remove this.
             prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
                 self.moe,
                 use_deepseek_fp8_block_scale=self.block_quant,

From 7d7d5a62256de6fe20e146ffd1187014246322cf Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 21:33:17 -0500
Subject: [PATCH 11/14] nit changes

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 vllm/model_executor/layers/fused_moe/all2all_utils.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py
index d76c629b72ba..d6740d8a3ff5 100644
--- a/vllm/model_executor/layers/fused_moe/all2all_utils.py
+++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py
@@ -176,5 +176,4 @@ def maybe_make_prepare_finalize(
     elif moe.use_naive_kernels:
         prepare_finalize = MoEPrepareAndFinalizeNaiveEP()
 
-    print(f"{prepare_finalize=}")
     return prepare_finalize

From 63357f7d881959a3b97c1dde7043b9d41695f7a4 Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 21:55:16 -0500
Subject: [PATCH 12/14] support apply router weight on input

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../layers/fused_moe/prepare_finalize.py      | 35 +++++++++++++++----
 1 file changed, 28 insertions(+), 7 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index f406df814eff..c8ab382ab158 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -40,13 +40,13 @@ def prepare(
         apply_router_weight_on_input: bool,
         quant_config: FusedMoEQuantConfig,
     ) -> mk.PrepareResultType:
-        a1, topk_weights, topk_ids = get_ep_group().dispatch(
-            a1,
-            topk_weights,
-            topk_ids,
-            # TODO?
-            is_sequence_parallel=False,
-        )
+        if apply_router_weight_on_input:
+            topk = topk_ids.size(1)
+            assert topk == 1, (
+                "apply_router_weight_on_input is only implemented for topk=1"
+            )
+            # Note: do not use inplace for shared experts overlap
+            a1 = a1 * topk_weights.to(a1.dtype)
 
         a1q, a1q_scale = moe_kernel_quantize_input(
             a1,
@@ -58,6 +58,27 @@ def prepare(
         # TODO - this is just for deepgemm?
         expert_tokens_meta = None
 
+        from vllm.platforms import current_platform
+
+        # The torch ops do not support fp8, so use an int8 view.
+        # Since this does not do a reduce, it is safe to do.
+        use_int8_view = a1q.dtype == current_platform.fp8_dtype()
+        if use_int8_view:
+            a1q = a1q.view(torch.int8)
+
+        extra_tensors = None if a1q_scale is None else [a1q_scale]
+        a1q, topk_weights, topk_ids, et = get_ep_group().dispatch(
+            a1q,
+            topk_weights,
+            topk_ids,
+            is_sequence_parallel=False,  # TODO: support SP
+            extra_tensors=extra_tensors,
+        )
+        a1q_scale = et[0] if extra_tensors is not None else None
+
+        if use_int8_view:
+            a1q = a1q.view(current_platform.fp8_dtype())
+
         return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights
 
     def finalize(

From 3886cfbb4414aefe70dede928467eabed934b688 Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 22:09:01 -0500
Subject: [PATCH 13/14] attempt to get everything working for llama scout
 modelopt flashinfer cutlass

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../flashinfer_cutlass_prepare_finalize.py    | 11 +++---
 .../layers/fused_moe/prepare_finalize.py      | 34 ++++++++++++-------
 vllm/model_executor/models/llama4.py          |  1 +
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
index 0b0efdafbd4d..5f48e83ee5a3 100644
--- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
@@ -11,6 +11,7 @@
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
+    MoEPrepareAndFinalizeNaiveEP,
     MoEPrepareAndFinalizeNoEP,
 )
 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -363,13 +364,11 @@ def create_flashinfer_prepare_finalize(
         else:
             return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)
 
-    # FP8 DP path currently supported via AllGather.
+    # NOTE(rob): CUTLASS FP8 block quant executes the input
+    # quantzation and grouped gemm in a single kernel.
     if use_dp:
-        return FlashInferAllGatherMoEPrepareAndFinalize(
-            use_dp=True,
-            use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
+        return MoEPrepareAndFinalizeNaiveEP(
+            defer_input_quant=use_deepseek_fp8_block_scale
         )
     else:
-        # NOTE(rob): CUTLASS FP8 block quant executes the input
-        # quantzation and grouped gemm in a single kernel.
         return MoEPrepareAndFinalizeNoEP(defer_input_quant=use_deepseek_fp8_block_scale)
diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index c8ab382ab158..a6d81ea16500 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -14,6 +14,10 @@
 
 
 class MoEPrepareAndFinalizeNaiveEP(mk.FusedMoEPrepareAndFinalize):
+    def __init__(self, defer_input_quant: bool = False) -> None:
+        super().__init__()
+        self.defer_input_quant = defer_input_quant
+
     @property
     def activation_format(self) -> mk.FusedMoEActivationFormat:
         return mk.FusedMoEActivationFormat.Standard
@@ -48,33 +52,38 @@ def prepare(
             # Note: do not use inplace for shared experts overlap
             a1 = a1 * topk_weights.to(a1.dtype)
 
-        a1q, a1q_scale = moe_kernel_quantize_input(
-            a1,
-            quant_config.a1_scale,
-            quant_config.quant_dtype,
-            quant_config.per_act_token_quant,
-            quant_config.block_shape,
-        )
+        # Defer input quantization to the MoE kernel.
+        if self.defer_input_quant:
+            a1q = a1
+            a1q_scale = None
+        else:
+            a1q, a1q_scale = moe_kernel_quantize_input(
+                a1,
+                quant_config.a1_scale,
+                quant_config.quant_dtype,
+                quant_config.per_act_token_quant,
+                quant_config.block_shape,
+            )
         # TODO - this is just for deepgemm?
         expert_tokens_meta = None
 
         from vllm.platforms import current_platform
 
         # The torch ops do not support fp8, so use an int8 view.
-        # Since this does not do a reduce, it is safe to do.
+        # Since dispatch does not do a reduce, this is safe to do.
         use_int8_view = a1q.dtype == current_platform.fp8_dtype()
         if use_int8_view:
             a1q = a1q.view(torch.int8)
 
-        extra_tensors = None if a1q_scale is None else [a1q_scale]
-        a1q, topk_weights, topk_ids, et = get_ep_group().dispatch(
+        scales = None if a1q_scale is None else [a1q_scale]
+        a1q, topk_weights, topk_ids, scales = get_ep_group().dispatch(
             a1q,
             topk_weights,
             topk_ids,
             is_sequence_parallel=False,  # TODO: support SP
-            extra_tensors=extra_tensors,
+            extra_tensors=scales,
         )
-        a1q_scale = et[0] if extra_tensors is not None else None
+        a1q_scale = scales[0] if scales is not None else None
 
         if use_int8_view:
             a1q = a1q.view(current_platform.fp8_dtype())
@@ -101,7 +110,6 @@ def finalize(
             apply_router_weight_on_input=apply_router_weight_on_input,
         )
 
-        # ... copy the
         output.copy_(get_ep_group().combine(out, is_sequence_parallel=False))
 
 
diff --git a/vllm/model_executor/models/llama4.py b/vllm/model_executor/models/llama4.py
index 9ed0741acba1..3698d867c5a5 100644
--- a/vllm/model_executor/models/llama4.py
+++ b/vllm/model_executor/models/llama4.py
@@ -432,6 +432,7 @@ def load_moe_expert_weights(
 
         # Whether the MoE expert weights are loaded successfully.
         expert_param_loaded = False
+        loaded_weight = loaded_weight.to("cuda")
 
         # If fused is True, the loaded weight is in the layout of:
         # [num_experts, hidden_in, hidden_out], so we must transpose the last

From 2284b59ee76960d9c2d5e6b66a97bb4ffa06fcdf Mon Sep 17 00:00:00 2001
From: Robert Shaw <robshaw@redhat.com>
Date: Wed, 7 Jan 2026 22:36:24 -0500
Subject: [PATCH 14/14] updated

Signed-off-by: Robert Shaw <robshaw@redhat.com>
---
 .../layers/fused_moe/prepare_finalize.py          | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize.py b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
index a6d81ea16500..b66e2a1dd73f 100644
--- a/vllm/model_executor/layers/fused_moe/prepare_finalize.py
+++ b/vllm/model_executor/layers/fused_moe/prepare_finalize.py
@@ -75,15 +75,24 @@ def prepare(
         if use_int8_view:
             a1q = a1q.view(torch.int8)
 
-        scales = None if a1q_scale is None else [a1q_scale]
-        a1q, topk_weights, topk_ids, scales = get_ep_group().dispatch(
+        # Skip gathering scales if we have static quantization
+        # (the scale is a scalar, replicated on all ranks) or
+        # if quantization is deferred.
+        skip_gather_scales = a1q_scale is None or a1q_scale.ndim == 0
+        scales = None if skip_gather_scales else [a1q_scale]
+
+        res = get_ep_group().dispatch(
             a1q,
             topk_weights,
             topk_ids,
             is_sequence_parallel=False,  # TODO: support SP
             extra_tensors=scales,
         )
-        a1q_scale = scales[0] if scales is not None else None
+        if skip_gather_scales:
+            a1q, topk_weights, topk_ids = res
+        else:
+            a1q, topk_weights, topk_ids, scales = res
+            a1q_scale = res[0]
 
         if use_int8_view:
             a1q = a1q.view(current_platform.fp8_dtype())