There appears to be an issue with how dimensions are calculated for the Marlin backend, which could lead to incorrect CUDA buffer allocations and potential runtime errors.

1. N is being assigned the packed dimension 2*N_packed from qinfo.w13_qweight.shape. However, it's used to allocate intermediate_cache1, which requires the full gate_up_dim. It should be multiplied by 16 to get the unpacked dimension.
2. hidden_dim is being assigned qinfo.w2_qweight.shape[1], which is the packed intermediate dimension (N_packed), not the hidden dimension.

To fix this, you should use the unpacked dimensions for allocation.

The calculation for N (intermediate dimension) seems incorrect. quant_info.w2_qweight has a shape of [E, N_packed, K], where E is number of experts, N_packed is the packed intermediate dimension, and K is the hidden dimension. Therefore, shape[1] corresponds to N_packed. The correct way to get the intermediate dimension N would be to use shape[2] from w13_qweight or w2_qweight and unpack it, or use w13_qweight.shape[1] and unpack it.

Given the shapes, quant_info.w2_qweight.shape[2] would give the hidden dimension, which is also not correct for N. The intermediate dimension N should be derived from w13_qweight's second dimension, unpacked.

For better type safety and maintainability, consider using the specific LoRAHooks type instead of Any for the hooks parameter. This makes the expected type explicit and improves code clarity.

You would need to add a forward reference import within the TYPE_CHECKING block:

if TYPE_CHECKING:
    from sglang.srt.lora.lora_moe_runners import LoRAHooks
    ...

And then change the signature to:

    def run(
        self,
        runner_input: RunnerInput,
        quant_info: MoeQuantInfo,
        running_state: dict,
        hooks: Optional["LoRAHooks"] = None,
    ) -> RunnerOutput:

For better type safety and maintainability, you can replace Any with a more specific type hint for the hooks parameter. Using "LoRAHooks" as a forward reference would make the code clearer about the expected object type.

This would involve adding an import inside the TYPE_CHECKING block:

if TYPE_CHECKING:
    from sglang.srt.lora.lora_moe_runners import LoRAHooks
    ...

And updating the method signature accordingly.

To improve type safety, consider replacing Any with the specific "LoRAHooks" type hint for the hooks parameter. This will make the expected interface for hooks explicit.

You'll need to add from sglang.srt.lora.lora_moe_runners import LoRAHooks inside the TYPE_CHECKING block at the top of the file.

For better type safety and code clarity, it's recommended to use the specific "LoRAHooks" type hint instead of Any for the hooks parameter, even if it's not used in this particular implementation. This maintains API consistency and improves readability.

You would need to add from sglang.srt.lora.lora_moe_runners import LoRAHooks to the TYPE_CHECKING block.

@klshuster @yushengsu-thu Do not define a function in the forward / run critical path. Clean this up!