sglang 0.5.4__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/layers/quantization/modelopt_quant.py

@@ -25,10 +25,11 @@ from sglang.srt.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
+from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.layers.quantization.fp8_utils import (
     apply_fp8_linear,
     cutlass_fp8_supported,
-    is_sm100_supported,
+    is_blackwell_supported,
 )
 from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
@@ -49,8 +50,10 @@ if TYPE_CHECKING:
     )
     from sglang.srt.single_batch_overlap import DownGemmOverlapArgs
 
-if is_cuda():
-    from sgl_kernel import scaled_fp4_quant
+try:
+    from flashinfer import fp4_quantize
+except ImportError:
+    fp4_quantize = None
 
 try:
     from flashinfer import mm_fp4 as fp4_gemm
@@ -466,8 +469,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
             # Fp8 moe kernel needs single weight scale for w13 per expert.
             # We take the max of the w1 and w3 scales then dequant and requant each expert.
             if layer.w13_weight_scale.dim() == 2:  # Shape: (num_experts, 2)
-                from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
-
                 # Get the maximum scale across w1 and w3 for each expert
                 max_w13_scales = layer.w13_weight_scale.max(dim=1).values
 
@@ -515,6 +516,84 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
                 layer.w2_input_scale.max(), requires_grad=False
             )
 
+        # Align FP8 weights to FlashInfer per-tensor kernel layout if enabled
+        if should_use_flashinfer_trtllm_moe():
+            from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
+
+            # 1) Swap W13 halves: [Up, Gate] -> [Gate, Up] expected by FI
+            num_experts, two_n, hidden = layer.w13_weight.shape
+            inter = two_n // 2
+            w13_swapped = (
+                layer.w13_weight.reshape(num_experts, 2, inter, hidden)
+                .flip(dims=[1])
+                .reshape(num_experts, two_n, hidden)
+            )
+
+            # 2) Reorder rows for fused gated activation (W13)
+            w13_interleaved = [
+                reorder_rows_for_gated_act_gemm(w13_swapped[i])
+                for i in range(num_experts)
+            ]
+            w13_interleaved = torch.stack(w13_interleaved).reshape(
+                num_experts, two_n, hidden
+            )
+
+            # 3) Shuffle weights for transposed MMA output (both W13, W2)
+            epilogue_tile_m = 128
+            w13_shuffled = [
+                shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m)
+                for i in range(num_experts)
+            ]
+            w2_shuffled = [
+                shuffle_matrix_a(layer.w2_weight[i].view(torch.uint8), epilogue_tile_m)
+                for i in range(num_experts)
+            ]
+
+            layer.w13_weight = Parameter(
+                torch.stack(w13_shuffled).view(torch.float8_e4m3fn),
+                requires_grad=False,
+            )
+            layer.w2_weight = Parameter(
+                torch.stack(w2_shuffled).view(torch.float8_e4m3fn),
+                requires_grad=False,
+            )
+
+        # Precompute and register per-expert output scaling factors for FI MoE
+        if should_use_flashinfer_trtllm_moe():
+            # Note: w13_input_scale and w2_input_scale are scalar Parameters post-reduction
+            assert (
+                hasattr(layer, "w13_input_scale") and layer.w13_input_scale is not None
+            )
+            assert hasattr(layer, "w2_input_scale") and layer.w2_input_scale is not None
+            assert (
+                hasattr(layer, "w13_weight_scale")
+                and layer.w13_weight_scale is not None
+            )
+            assert (
+                hasattr(layer, "w2_weight_scale") and layer.w2_weight_scale is not None
+            )
+
+            input_scale = layer.w13_input_scale.to(torch.float32)
+            activation_scale = layer.w2_input_scale.to(torch.float32)
+            w13_weight_scale = layer.w13_weight_scale.to(torch.float32)
+            w2_weight_scale = layer.w2_weight_scale.to(torch.float32)
+
+            output1_scales_scalar = (
+                w13_weight_scale * input_scale * (1.0 / activation_scale)
+            )
+            output1_scales_gate_scalar = w13_weight_scale * input_scale
+            output2_scales_scalar = activation_scale * w2_weight_scale
+
+            layer.output1_scales_scalar = Parameter(
+                output1_scales_scalar, requires_grad=False
+            )
+            layer.output1_scales_gate_scalar = Parameter(
+                output1_scales_gate_scalar, requires_grad=False
+            )
+            layer.output2_scales_scalar = Parameter(
+                output2_scales_scalar, requires_grad=False
+            )
+
     def create_moe_runner(
         self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
     ):
@@ -526,6 +605,81 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+
+        # Fast path: TRT-LLM FP8 per-tensor MoE using BYPASSED TopK routing
+        from sglang.srt.layers.moe.topk import TopKOutputChecker
+
+        if should_use_flashinfer_trtllm_moe() and TopKOutputChecker.format_is_bypassed(
+            topk_output
+        ):
+            router_logits = topk_output.router_logits
+            topk_config = topk_output.topk_config
+
+            # Constraints
+            assert (
+                self.moe_runner_config.activation == "silu"
+            ), "Only silu is supported for flashinfer fp8 moe"
+
+            from flashinfer import RoutingMethodType
+            from flashinfer.fused_moe import trtllm_fp8_per_tensor_scale_moe
+
+            correction_bias = (
+                None
+                if topk_config.correction_bias is None
+                else topk_config.correction_bias
+            )
+            # Pre-quantize activations to FP8 per-tensor using provided input scale
+            x_fp8, _ = scaled_fp8_quant(x, layer.w13_input_scale)
+
+            use_routing_scales_on_input = True
+            routed_scaling_factor = self.moe_runner_config.routed_scaling_factor
+
+            # Enforce Llama4 routing for ModelOpt FP8 MoE for now.
+            # TODO(brayden): support other routing methods
+            assert topk_config.top_k == 1, "ModelOpt FP8 MoE requires top_k==1"
+            assert (
+                not topk_config.num_expert_group
+            ), "ModelOpt FP8 MoE does not support expert grouping"
+            assert (
+                not topk_config.topk_group
+            ), "ModelOpt FP8 MoE does not support grouped top-k"
+            routing_method_type = RoutingMethodType.Llama4
+
+            # FlashInfer TRTLLM requires routing_logits (and bias) to be bfloat16
+            routing_logits_cast = router_logits.to(torch.bfloat16)
+            routing_bias_cast = (
+                None if correction_bias is None else correction_bias.to(torch.bfloat16)
+            )
+
+            output = trtllm_fp8_per_tensor_scale_moe(
+                routing_logits=routing_logits_cast,
+                routing_bias=routing_bias_cast,
+                hidden_states=x_fp8,
+                gemm1_weights=layer.w13_weight,
+                output1_scales_scalar=layer.output1_scales_scalar,
+                output1_scales_gate_scalar=layer.output1_scales_gate_scalar,
+                gemm2_weights=layer.w2_weight,
+                output2_scales_scalar=layer.output2_scales_scalar,
+                num_experts=layer.num_experts,
+                top_k=topk_config.top_k,
+                n_group=0,
+                topk_group=0,
+                intermediate_size=layer.w2_weight.shape[2],
+                local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
+                local_num_experts=layer.num_local_experts,
+                routed_scaling_factor=(
+                    routed_scaling_factor if routed_scaling_factor is not None else 1.0
+                ),
+                use_routing_scales_on_input=use_routing_scales_on_input,
+                tile_tokens_dim=8,  # TODO(brayden): use the FI tile calculation
+                routing_method_type=routing_method_type,
+            )
+
+            from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+            return StandardCombineInput(hidden_states=output)
 
         quant_info = TritonMoeQuantInfo(
             w13_weight=layer.w13_weight,
@@ -867,10 +1021,9 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         output_shape = [x_m, w_n]
 
         # Quantize BF16 or FP16 to (FP4 and interleaved block scale)
-        x_fp4, x_scale_interleaved = scaled_fp4_quant(x, layer.input_scale_inv)
+        x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
 
         assert x_fp4.dtype == torch.uint8
-        assert x_scale_interleaved.dtype == torch.float8_e4m3fn
         assert layer.weight.dtype == torch.uint8
         assert layer.weight_scale_interleaved.dtype == torch.float8_e4m3fn
         assert layer.alpha.dtype == torch.float32
@@ -903,7 +1056,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
 
     def __init__(self, quant_config: ModelOptFp4Config):
         self.quant_config = quant_config
-        if not is_sm100_supported():
+        if not is_blackwell_supported():
             raise ValueError(
                 "Current platform does not support NVFP4"
                 " quantization. Please use Blackwell and"
@@ -1383,8 +1536,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         alt_stream=None,
     ) -> CombineInput:
 
-        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
-
         x = dispatch_output.hidden_states
         topk_output = dispatch_output.topk_output
 
@@ -1397,6 +1548,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
         # Check if this is a FlashInferFP4MoE layer that should handle its own forward
         if hasattr(layer, "gemm1_weights_fp4_shuffled"):
             # This layer was processed with flashinfer TRTLLM - delegate to its own forward
+            from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
             return StandardCombineInput(hidden_states=layer.forward(x, topk_output))
 
         if self.enable_flashinfer_cutlass_moe:
@@ -1410,7 +1563,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
             output_dtype = x.dtype
             x_sf = None
             if should_use_flashinfer_cutlass_moe_fp4_allgather():
-                from flashinfer import fp4_quantize, nvfp4_block_scale_interleave
+                from flashinfer import nvfp4_block_scale_interleave
 
                 # Quantize before comm, swizzle after.
                 if x.shape[0] > 0:
@@ -1465,6 +1618,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
                 if forward_shared_experts is not None:
                     torch.cuda.current_stream().wait_stream(alt_stream)
 
+            from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
             return StandardCombineInput(hidden_states=output)
 
         from sglang.srt.layers.moe.cutlass_moe import cutlass_moe_fp4
@@ -1486,6 +1641,8 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
             apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input,
         ).to(x.dtype)
         # Scale by routed_scaling_factor is fused into select_experts.
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
         return StandardCombineInput(hidden_states=output)
 
     def apply_without_routing_weights(

sglang/srt/layers/quantization/mxfp4.py

@@ -261,26 +261,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
 
         self.prefix = prefix
         self.topk_indices_dtype = None
-        self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel()
+        self.use_triton_kernels = get_moe_runner_backend().is_triton_kernels()
         self.with_bias = False
         self.use_flashinfer = get_moe_runner_backend().is_flashinfer_mxfp4()
         self.flashinfer_mxfp4_moe_precision = (
             get_global_server_args().flashinfer_mxfp4_moe_precision
         )
 
-        self.triton_kernel_moe_forward = None
-        self.triton_kernel_moe_with_bias_forward = None
-        if torch.cuda.is_available() and has_triton_kernels:
-            from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
-                triton_kernel_moe_forward as _tk_forward,
-            )
-            from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
-                triton_kernel_moe_with_bias_forward as _tk_with_bias_forward,
-            )
-
-            self.triton_kernel_moe_forward = _tk_forward
-            self.triton_kernel_moe_with_bias_forward = _tk_with_bias_forward
-
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -600,7 +587,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
         self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
     ):
         self.moe_runner_config = moe_runner_config
-        self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
+        backend = (
+            MoeRunnerBackend.TRITON_KERNELS
+            if self.use_triton_kernels
+            else MoeRunnerBackend.TRITON
+        )
+        self.runner = MoeRunner(backend, moe_runner_config)
 
     def apply(
         self,
@@ -677,31 +669,31 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
             )[0]
             return StandardCombineInput(hidden_states=trtllm_gen_output)
 
-        if self.use_triton_kernels:
+        backend = self.runner.runner_backend
+        if backend.is_triton_kernels():
+            from sglang.srt.layers.moe.moe_runner.triton_kernels import (
+                TritonKernelsQuantInfo,
+            )
+
             assert (
                 layer.moe_ep_size == 1
             ), "Expert parallel is not supported when using triton kernels"
-            if self.with_bias:
-                output = self.triton_kernel_moe_with_bias_forward(
-                    hidden_states=x,
-                    w1=self.w13_weight_triton_tensor,
-                    w1_pcg=self.w13_precision_config,
-                    w2=self.w2_weight_triton_tensor,
-                    w2_pcg=self.w2_precision_config,
-                    b1=layer.w13_weight_bias,
-                    b2=layer.w2_weight_bias,
-                    topk_output=topk_output,
-                    moe_runner_config=moe_runner_config,
-                )
-            else:
-                output = self.triton_kernel_moe_forward(
-                    hidden_states=x,
-                    w1=layer.w13_weight,
-                    w2=layer.w2_weight,
-                    topk_output=topk_output,
-                    moe_runner_config=moe_runner_config,
-                )
-            return StandardCombineInput(hidden_states=output)
+            quant_info = TritonKernelsQuantInfo(
+                w13_weight=(
+                    self.w13_weight_triton_tensor
+                    if self.w13_weight_triton_tensor is not None
+                    else layer.w13_weight
+                ),
+                w2_weight=(
+                    self.w2_weight_triton_tensor
+                    if self.w2_weight_triton_tensor is not None
+                    else layer.w2_weight
+                ),
+                w13_bias=getattr(layer, "w13_weight_bias", None),
+                w2_bias=getattr(layer, "w2_weight_bias", None),
+                w13_precision_config=getattr(self, "w13_precision_config", None),
+                w2_precision_config=getattr(self, "w2_precision_config", None),
+            )
         else:
             quant_info = TritonMoeQuantInfo(
                 w13_weight=layer.w13_weight,
@@ -709,7 +701,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 b13=getattr(layer, "w13_weight_bias", None),
                 b2=getattr(layer, "w2_weight_bias", None),
             )
-            return self.runner.run(dispatch_output, quant_info)
+        return self.runner.run(dispatch_output, quant_info)
 
 
 class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase):

sglang/srt/layers/quantization/unquant.py

@@ -115,13 +115,15 @@ class UnquantizedLinearMethod(LinearMethodBase):
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-
         if use_intel_amx_backend(layer):
             x_shapes = x.shape
             if len(x_shapes) == 3:
                 x = x.view(-1, x.shape[-1])
             output = torch.ops.sgl_kernel.weight_packed_linear(
-                x, layer.weight, bias, True  # is_vnni
+                x,
+                layer.weight,
+                bias,
+                True,  # is_vnni
             )
             if len(x_shapes) == 3:
                 output = output.view(x_shapes[0], x_shapes[1], -1)
@@ -138,19 +140,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         self.use_triton_kernels = use_triton_kernels
         self.with_bias = False
 
-        self.triton_kernel_moe_forward = None
-        self.triton_kernel_moe_with_bias_forward = None
-        if torch.cuda.is_available() and use_triton_kernels:
-            from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
-                triton_kernel_moe_forward as _tk_forward,
-            )
-            from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
-                triton_kernel_moe_with_bias_forward as _tk_with_bias_forward,
-            )
-
-            self.triton_kernel_moe_forward = _tk_forward
-            self.triton_kernel_moe_with_bias_forward = _tk_with_bias_forward
-
     def create_weights(
         self,
         layer: torch.nn.Module,
@@ -231,14 +220,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
     ):
         self.moe_runner_config = moe_runner_config
-        self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
+        backend = (
+            MoeRunnerBackend.TRITON_KERNELS
+            if self.use_triton_kernels
+            else MoeRunnerBackend.TRITON
+        )
+        self.runner = MoeRunner(backend, moe_runner_config)
 
     def apply(
         self,
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
-
         return self.forward(
             layer=layer,
             dispatch_output=dispatch_output,
@@ -249,7 +242,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
-
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
 
         x = dispatch_output.hidden_states
@@ -257,30 +249,19 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
 
         moe_runner_config = self.moe_runner_config
 
-        if self.use_triton_kernels:
-            if self.with_bias:
-                assert self.triton_kernel_moe_with_bias_forward is not None
-                output = self.triton_kernel_moe_with_bias_forward(
-                    hidden_states=x,
-                    w1=layer.w13_weight,
-                    w2=layer.w2_weight,
-                    b1=layer.w13_weight_bias,
-                    b2=layer.w2_weight_bias,
-                    topk_output=topk_output,
-                    moe_runner_config=moe_runner_config,
-                    w1_pcg=None,
-                    w2_pcg=None,
-                )
-            else:
-                assert self.triton_kernel_moe_forward is not None
-                output = self.triton_kernel_moe_forward(
-                    hidden_states=x,
-                    w1=layer.w13_weight,
-                    w2=layer.w2_weight,
-                    topk_output=topk_output,
-                    moe_runner_config=moe_runner_config,
-                )
-            return StandardCombineInput(hidden_states=output)
+        backend = self.runner.runner_backend
+        if backend.is_triton_kernels():
+            from sglang.srt.layers.moe.moe_runner.triton_kernels import (
+                TritonKernelsQuantInfo,
+            )
+
+            quant_info = TritonKernelsQuantInfo(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                w13_bias=getattr(layer, "w13_weight_bias", None),
+                w2_bias=getattr(layer, "w2_weight_bias", None),
+            )
+            return self.runner.run(dispatch_output, quant_info)
         else:
             if _use_aiter:
                 assert not moe_runner_config.no_combine, "unsupported"
@@ -311,7 +292,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
                 )
                 return StandardCombineInput(hidden_states=output)
             else:
-
                 quant_info = TritonMoeQuantInfo(
                     w13_weight=layer.w13_weight,
                     w2_weight=layer.w2_weight,
@@ -325,7 +305,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
-
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
 
         x = dispatch_output.hidden_states
@@ -380,7 +359,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
         layer: torch.nn.Module,
         dispatch_output: StandardDispatchOutput,
     ) -> CombineInput:
-
         import torch_npu
 
         from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput

sglang/srt/layers/quantization/w4afp8.py

@@ -23,7 +23,8 @@ if TYPE_CHECKING:
     from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
     from sglang.srt.layers.moe.token_dispatcher import (
         CombineInput,
-        DeepEPNormalOutput,
+        DeepEPLLDispatchOutput,
+        DeepEPNormalDispatchOutput,
         StandardDispatchOutput,
     )
 
@@ -328,10 +329,45 @@ class W4AFp8MoEMethod(FusedMoEMethodBase):
             output *= self.moe_runner_config.routed_scaling_factor
         return StandardCombineInput(hidden_states=output)
 
+    def apply_deepep_ll(
+        self,
+        layer: DeepEPMoE,
+        dispatch_output: DeepEPLLDispatchOutput,
+    ) -> torch.Tensor:
+
+        from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe_deepep_ll
+
+        hidden_states, _, topk_ids, _, masked_m, _ = dispatch_output
+
+        output = cutlass_w4a8_moe_deepep_ll(
+            hidden_states,
+            layer.w13_weight,
+            layer.w2_weight,
+            layer.w13_weight_scale_inv,
+            layer.w2_weight_scale_inv,
+            topk_ids,
+            masked_m,
+            layer.quant_method.a_strides1,
+            layer.quant_method.b_strides1,
+            layer.quant_method.c_strides1,
+            layer.quant_method.a_strides2,
+            layer.quant_method.b_strides2,
+            layer.quant_method.c_strides2,
+            layer.quant_method.s_strides13,
+            layer.quant_method.s_strides2,
+            layer.quant_method.expert_offsets,
+            layer.quant_method.problem_sizes1,
+            layer.quant_method.problem_sizes2,
+            layer.w13_input_scale,
+            layer.w2_input_scale,
+        )
+
+        return output
+
     def apply_deepep_normal(
         self,
         layer: DeepEPMoE,
-        dispatch_output: DeepEPNormalOutput,
+        dispatch_output: DeepEPNormalDispatchOutput,
     ) -> torch.Tensor:
         from sglang.srt.layers.moe.cutlass_w4a8_moe import (
             cutlass_w4a8_moe_deepep_normal,

sglang/srt/layers/radix_attention.py

@@ -142,8 +142,11 @@ def unified_attention_with_output(
     ret = forward_batch.attn_backend.forward(
         query, key, value, attention_layer, forward_batch, save_kv_cache
     )
-    assert output.shape == ret.shape
-    output.copy_(ret)
+    assert (
+        output.numel() == ret.numel()
+    ), f"Output tensor element mismatch: {output.numel()} != {ret.numel()}"
+
+    output.view(ret.shape).copy_(ret)
     return