From 9a190f142c8acccffd0b19d2393e1eb76f775aea Mon Sep 17 00:00:00 2001
From: Ashraf Bhuiyan <mbhuiyan@redhat.com>
Date: Tue, 29 Sep 2026 14:13:52 -0400
Subject: [PATCH] [Mypy] Fix mypy typing for Zamba2 models (#58255)

Signed-off-by: Ashraf Bhuiyan <mbhuiyan@redhat.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
---
 tools/pre_commit/mypy.py             |  1 -
 vllm/model_executor/models/zamba2.py | 25 +++++++++++++++----------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/tools/pre_commit/mypy.py b/tools/pre_commit/mypy.py
index 5ad37e138e..88015c6a53 100755
--- a/tools/pre_commit/mypy.py
+++ b/tools/pre_commit/mypy.py
@@ -113,7 +113,6 @@ SEPARATE_GROUPS = [
 
 EXCLUDE = [
     r"vllm/model_executor/models/[kK]",
-    r"vllm/model_executor/models/[zZ]",
 ]
 
 
diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py
index 9398d99768..6506c0776e 100644
--- a/vllm/model_executor/models/zamba2.py
+++ b/vllm/model_executor/models/zamba2.py
@@ -86,16 +86,21 @@ class Zamba2LoRA(nn.Module):
         )
 
         if isinstance(output_dim, list):
-            B_class = MergedColumnParallelLinear
+            self.B: ColumnParallelLinear = MergedColumnParallelLinear(
+                rank,
+                output_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.B",
+            )
         else:
-            B_class = ColumnParallelLinear
-        self.B = B_class(
-            rank,
-            output_dim,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.B",
-        )
+            self.B = ColumnParallelLinear(
+                rank,
+                output_dim,
+                bias=False,
+                quant_config=quant_config,
+                prefix=f"{prefix}.B",
+            )
 
     def forward(
         self,
@@ -304,7 +309,7 @@ class Zamba2MLP(nn.Module):
         self,
         config: Zamba2Config,
         bare_block_idx: int,
-        num_hybrid_layers: dict[int, int],
+        num_hybrid_layers: int,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None: