Improve LoRA Perf by Deprecating FlashInfer and Eliminating Redundant Tensor Ops (#8940)

python/sglang/srt/models/gemma3n_mm.py

@@ -501,23 +501,16 @@ class Gemma3nForConditionalGeneration(PreTrainedModel):
 
     def get_hidden_dim(self, module_name):
         # return input_dim, output_dim
-        # TODO: the special handling of qkv will be addressed in #8940.
         if module_name == "qkv_proj":
             return (
                 self.config.hidden_size,
-                None,  # qkv_proj is only used in LoRA A
+                self.config.head_dim
+                * (
+                    self.config.num_attention_heads
+                    + self.config.num_key_value_heads * 2
+                ),
             )
-        elif module_name == "kv_proj":
-            return (
-                None,  # kv_proj is only used in LoRA B
-                self.config.head_dim * self.config.num_key_value_heads,
-            )
-        elif module_name == "q_proj":
-            return (
-                None,  # q_proj is only used in LoRA B
-                self.config.head_dim * self.config.num_attention_heads,
-            )
-        elif module_name in ["o_proj"]:
+        elif module_name == "o_proj":
             return (
                 self.config.head_dim * self.config.num_attention_heads,
                 self.config.hidden_size,
@@ -527,7 +520,7 @@ class Gemma3nForConditionalGeneration(PreTrainedModel):
                 "Currently SGLang requires uniform intermediate size for all layers. "
                 "Please file an issue if you need support for non-uniform intermediate sizes."
             )
-            return self.config.hidden_size, self.config.intermediate_size[0]
+            return self.config.hidden_size, self.config.intermediate_size[0] * 2
         elif module_name == "down_proj":
             assert len(set(self.config.intermediate_size)) == 1, (
                 "Currently SGLang requires uniform intermediate size for all layers. "