[LoRA] Laguna: per-layer LoRA hidden-dim resolution for packed attention (#30298)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/lora/utils.py

@@ -124,106 +124,114 @@ def get_hidden_dim(
         Please implement the function in the model class if it is not.
         You can reference this function in llama.py.
         """
-        head_dim = getattr(
-            config, "head_dim", config.hidden_size // config.num_attention_heads
+        return get_default_hidden_dim(
+            module_name, config, layer_idx, lora_added_vocab_size
         )
-        if module_name == "qkv_proj":
-            return config.hidden_size, head_dim * (
-                config.num_attention_heads + config.num_key_value_heads * 2
-            )
-        elif module_name == "o_proj":
-            o_head_dim = getattr(config, "v_head_dim", None) or head_dim
-            return (
-                o_head_dim * config.num_attention_heads,
-                config.hidden_size,
-            )
-        elif module_name == "gate_up_proj":
-            inter = config.intermediate_size
-            first_k = getattr(config, "first_k_dense_replace", None)
-            moe_freq = getattr(config, "moe_layer_freq", 1)
-            if (
-                first_k is not None
-                and layer_idx >= first_k
-                and layer_idx % moe_freq == 0
-            ):
-                moe_inter = getattr(config, "moe_intermediate_size", None)
-                n_shared = getattr(config, "n_shared_experts", None)
-                if moe_inter is not None and n_shared is not None:
-                    inter = moe_inter * n_shared
-            return config.hidden_size, inter * 2
-        elif module_name == "down_proj":
-            inter = config.intermediate_size
-            first_k = getattr(config, "first_k_dense_replace", None)
-            moe_freq = getattr(config, "moe_layer_freq", 1)
-            if (
-                first_k is not None
-                and layer_idx >= first_k
-                and layer_idx % moe_freq == 0
-            ):
-                moe_inter = getattr(config, "moe_intermediate_size", None)
-                n_shared = getattr(config, "n_shared_experts", None)
-                if moe_inter is not None and n_shared is not None:
-                    inter = moe_inter * n_shared
-            return inter, config.hidden_size
-        elif module_name == "fused_qkv_a_proj_with_mqa":
-            q_lora_rank = getattr(config, "q_lora_rank", None) or 0
-            kv_lora_rank = config.kv_lora_rank
-            qk_rope_head_dim = config.qk_rope_head_dim
-            return (
-                config.hidden_size,
-                q_lora_rank + kv_lora_rank + qk_rope_head_dim,
-            )
-        elif module_name == "q_b_proj":
+
+
+def get_default_hidden_dim(
+    module_name: str,
+    config: AutoConfig,
+    layer_idx: int,
+    lora_added_vocab_size: int = 0,
+) -> Tuple[int]:
+    """
+    Config-driven LoRA input/output dims for a module, assuming uniform
+    attention geometry across layers.
+
+    This is the fallback used when a model does not define ``get_hidden_dim``.
+    Models with per-layer geometry (e.g. Laguna's per-layer attention head
+    counts) should define ``get_hidden_dim`` on the model class, override the
+    layer-dependent modules there, and delegate the rest back to this helper
+    rather than re-deriving every branch.
+    """
+    head_dim = getattr(
+        config, "head_dim", config.hidden_size // config.num_attention_heads
+    )
+    if module_name == "qkv_proj":
+        return config.hidden_size, head_dim * (
+            config.num_attention_heads + config.num_key_value_heads * 2
+        )
+    elif module_name == "o_proj":
+        o_head_dim = getattr(config, "v_head_dim", None) or head_dim
+        return (
+            o_head_dim * config.num_attention_heads,
+            config.hidden_size,
+        )
+    elif module_name == "gate_up_proj":
+        inter = config.intermediate_size
+        first_k = getattr(config, "first_k_dense_replace", None)
+        moe_freq = getattr(config, "moe_layer_freq", 1)
+        if first_k is not None and layer_idx >= first_k and layer_idx % moe_freq == 0:
+            moe_inter = getattr(config, "moe_intermediate_size", None)
+            n_shared = getattr(config, "n_shared_experts", None)
+            if moe_inter is not None and n_shared is not None:
+                inter = moe_inter * n_shared
+        return config.hidden_size, inter * 2
+    elif module_name == "down_proj":
+        inter = config.intermediate_size
+        first_k = getattr(config, "first_k_dense_replace", None)
+        moe_freq = getattr(config, "moe_layer_freq", 1)
+        if first_k is not None and layer_idx >= first_k and layer_idx % moe_freq == 0:
+            moe_inter = getattr(config, "moe_intermediate_size", None)
+            n_shared = getattr(config, "n_shared_experts", None)
+            if moe_inter is not None and n_shared is not None:
+                inter = moe_inter * n_shared
+        return inter, config.hidden_size
+    elif module_name == "fused_qkv_a_proj_with_mqa":
+        q_lora_rank = getattr(config, "q_lora_rank", None) or 0
+        kv_lora_rank = config.kv_lora_rank
+        qk_rope_head_dim = config.qk_rope_head_dim
+        return (
+            config.hidden_size,
+            q_lora_rank + kv_lora_rank + qk_rope_head_dim,
+        )
+    elif module_name == "q_b_proj":
+        return (
+            config.q_lora_rank,
+            config.num_attention_heads
+            * (config.qk_nope_head_dim + config.qk_rope_head_dim),
+        )
+    elif module_name == "kv_b_proj":
+        return (
+            config.kv_lora_rank,
+            config.num_attention_heads * (config.qk_nope_head_dim + config.v_head_dim),
+        )
+    elif module_name in DSA_INDEXER_LORA_NAMES:
+        from sglang.srt.configs.model_config import (
+            get_dsa_index_head_dim,
+            get_dsa_index_n_heads,
+        )
+
+        if module_name == "indexer.wq_b":
             return (
                 config.q_lora_rank,
-                config.num_attention_heads
-                * (config.qk_nope_head_dim + config.qk_rope_head_dim),
-            )
-        elif module_name == "kv_b_proj":
-            return (
-                config.kv_lora_rank,
-                config.num_attention_heads
-                * (config.qk_nope_head_dim + config.v_head_dim),
-            )
-        elif module_name in DSA_INDEXER_LORA_NAMES:
-            from sglang.srt.configs.model_config import (
-                get_dsa_index_head_dim,
-                get_dsa_index_n_heads,
-            )
-
-            if module_name == "indexer.wq_b":
-                return (
-                    config.q_lora_rank,
-                    get_dsa_index_n_heads(config) * get_dsa_index_head_dim(config),
-                )
-            elif module_name == "indexer.wk":
-                return config.hidden_size, get_dsa_index_head_dim(config)
-            else:  # indexer.weights_proj
-                return config.hidden_size, get_dsa_index_n_heads(config)
-        elif module_name == "gate_up_proj_moe":
-            moe_inter = (
-                getattr(config, "moe_intermediate_size", None)
-                or config.intermediate_size
-            )
-            return config.hidden_size, moe_inter * 2
-        elif module_name == "down_proj_moe":
-            moe_inter = (
-                getattr(config, "moe_intermediate_size", None)
-                or config.intermediate_size
-            )
-            return moe_inter, config.hidden_size
-        elif module_name == "embed_tokens":
-            # For embedding: input is vocab_size (as embedding lookup), output is hidden_size
-            # if contain extra tokens will be added; otherwise is 0.
-            return config.vocab_size + lora_added_vocab_size, config.hidden_size
-        elif module_name == "lm_head":
-            # For lm_head: input is hidden_size, output is vocab_size
-            # if contain extra tokens will be added; otherwise is 0.
-            return config.hidden_size, config.vocab_size + lora_added_vocab_size
-        else:
-            raise NotImplementedError(
-                "get_hidden_dim not implemented for " + module_name
+                get_dsa_index_n_heads(config) * get_dsa_index_head_dim(config),
             )
+        elif module_name == "indexer.wk":
+            return config.hidden_size, get_dsa_index_head_dim(config)
+        else:  # indexer.weights_proj
+            return config.hidden_size, get_dsa_index_n_heads(config)
+    elif module_name == "gate_up_proj_moe":
+        moe_inter = (
+            getattr(config, "moe_intermediate_size", None) or config.intermediate_size
+        )
+        return config.hidden_size, moe_inter * 2
+    elif module_name == "down_proj_moe":
+        moe_inter = (
+            getattr(config, "moe_intermediate_size", None) or config.intermediate_size
+        )
+        return moe_inter, config.hidden_size
+    elif module_name == "embed_tokens":
+        # For embedding: input is vocab_size (as embedding lookup), output is hidden_size
+        # if contain extra tokens will be added; otherwise is 0.
+        return config.vocab_size + lora_added_vocab_size, config.hidden_size
+    elif module_name == "lm_head":
+        # For lm_head: input is hidden_size, output is vocab_size
+        # if contain extra tokens will be added; otherwise is 0.
+        return config.hidden_size, config.vocab_size + lora_added_vocab_size
+    else:
+        raise NotImplementedError("get_hidden_dim not implemented for " + module_name)
 
 
 def get_normalized_target_modules(

python/sglang/srt/models/laguna.py

@@ -50,6 +50,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from sglang.srt.lora.utils import get_default_hidden_dim
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.srt.models.utils import apply_qk_norm
@@ -566,6 +567,31 @@ class LagunaModel(nn.Module):
     def get_input_embeddings(self) -> nn.Embedding:
         return self.embed_tokens
 
+    def get_hidden_dim(self, module_name: str, layer_idx: int) -> Tuple[int, int]:
+        """LoRA input/output dims for a module, honoring Laguna's per-layer
+        attention widths.
+
+        Laguna sizes each layer's attention from
+        ``num_attention_heads_per_layer[layer_idx]`` (see ``LagunaAttention``),
+        so ``config.num_attention_heads`` — a single global value — is wrong for
+        any layer with a different head count. The generic
+        :func:`get_default_hidden_dim` fallback would use that global value and
+        mis-size the ``qkv_proj`` / ``o_proj`` LoRA buffers, crashing at
+        generation with ``sgemm_lora_a.py: assert x.shape[-1] == K``. We
+        override just those two attention projections and delegate every other
+        module (MLP, MoE, embed, lm_head, ...) to the shared helper.
+        """
+        config = self.config
+        # No fallback; Laguna's head_dim is non-standard.
+        head_dim = config.head_dim
+        num_heads = config.num_attention_heads_per_layer[layer_idx]
+        num_kv_heads = config.num_key_value_heads
+        if module_name == "qkv_proj":
+            return config.hidden_size, head_dim * (num_heads + num_kv_heads * 2)
+        elif module_name == "o_proj":
+            return head_dim * num_heads, config.hidden_size
+        return get_default_hidden_dim(module_name, config, layer_idx)
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -696,6 +722,9 @@ class LagunaForCausalLM(nn.Module):
     def get_input_embeddings(self) -> nn.Embedding:
         return self.model.embed_tokens
 
+    def get_hidden_dim(self, module_name: str, layer_idx: int) -> Tuple[int, int]:
+        return self.model.get_hidden_dim(module_name, layer_idx)
+
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
             ("qkv_proj", "q_proj", "q"),

test/registered/unit/lora/test_laguna_hidden_dim_unit.py

@@ -0,0 +1,163 @@
+"""Unit tests for Laguna's per-layer LoRA hidden-dim resolution.
+
+Laguna (`poolside/Laguna-XS.2`) sizes each layer's attention from
+`num_attention_heads_per_layer[layer_idx]`, so `config.num_attention_heads` —
+a single global value — is wrong for any layer whose head count differs. The
+generic `get_default_hidden_dim` fallback would use that global value and
+mis-size the `qkv_proj` / `o_proj` LoRA buffers, which crashes at generation
+with `sgemm_lora_a.py: assert x.shape[-1] == K`.
+
+`LagunaModel.get_hidden_dim` overrides just the two attention projections with
+per-layer widths and delegates every other module to the shared helper. These
+tests exercise that method directly against a minimal fake config — no CUDA,
+no server, no real weights — so they stay hermetic and fast.
+
+Usage:
+    python -m pytest test/registered/unit/lora/test_laguna_hidden_dim_unit.py -v
+"""
+
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+# CPU-only unit test; no CUDA/distributed dependencies.
+register_cuda_ci(est_time=6, stage="base-b", runner_config="1-gpu-small")
+register_amd_ci(est_time=6, suite="stage-b-test-1-gpu-small-amd")
+
+import unittest
+
+from sglang.srt.configs.laguna import LagunaConfig
+from sglang.srt.lora.utils import get_default_hidden_dim
+from sglang.srt.models.laguna import LagunaForCausalLM, LagunaModel
+
+
+def _make_fake_laguna(num_attention_heads_per_layer):
+    """Build a `LagunaModel` stand-in around a real `LagunaConfig`, without
+    running `__init__` (no weights, CPU-only).
+
+    `head_dim` (128) is deliberately not `hidden_size // num_attention_heads`,
+    matching every real Laguna checkpoint — so the hook must read `head_dim`
+    directly, never derive it. `LagunaConfig` sets the global
+    `num_attention_heads` from the first full-attention layer.
+    """
+    config = LagunaConfig(
+        hidden_size=2048,
+        head_dim=128,
+        num_key_value_heads=8,
+        num_hidden_layers=len(num_attention_heads_per_layer),
+        num_attention_heads_per_layer=list(num_attention_heads_per_layer),
+    )
+    model = LagunaModel.__new__(LagunaModel)
+    model.config = config
+    return model
+
+
+class TestLagunaPerLayerAttentionDims(unittest.TestCase):
+    """`qkv_proj` / `o_proj` must follow the *layer's own* head count."""
+
+    def setUp(self):
+        # Layer 0: 48 heads (== global default). Layer 1: 64 heads (differs).
+        self.model = _make_fake_laguna([48, 64])
+        self.head_dim = 128
+        self.hidden_size = 2048
+        self.num_kv_heads = 8
+
+    def test_qkv_proj_uses_first_layer_head_count(self):
+        got = self.model.get_hidden_dim("qkv_proj", layer_idx=0)
+        expected_out = self.head_dim * (48 + self.num_kv_heads * 2)
+        self.assertEqual(got, (self.hidden_size, expected_out))
+
+    def test_qkv_proj_uses_wider_layer_head_count(self):
+        # The layer that the global fallback would size incorrectly.
+        got = self.model.get_hidden_dim("qkv_proj", layer_idx=1)
+        expected_out = self.head_dim * (64 + self.num_kv_heads * 2)
+        self.assertEqual(got, (self.hidden_size, expected_out))
+
+    def test_o_proj_uses_first_layer_head_count(self):
+        got = self.model.get_hidden_dim("o_proj", layer_idx=0)
+        self.assertEqual(got, (self.head_dim * 48, self.hidden_size))
+
+    def test_o_proj_uses_wider_layer_head_count(self):
+        got = self.model.get_hidden_dim("o_proj", layer_idx=1)
+        self.assertEqual(got, (self.head_dim * 64, self.hidden_size))
+
+    def test_wider_layer_differs_from_generic_fallback(self):
+        """The regression guard: for the wider layer, the model hook must
+        return a DIFFERENT (correct) dim than the generic global-head
+        fallback — otherwise the buffer is mis-sized and generation asserts
+        in `sgemm_lora_a.py`.
+        """
+        for module_name in ("qkv_proj", "o_proj"):
+            hook = self.model.get_hidden_dim(module_name, layer_idx=1)
+            generic = get_default_hidden_dim(module_name, self.model.config, 1)
+            self.assertNotEqual(
+                hook,
+                generic,
+                f"{module_name}: per-layer hook must diverge from the global "
+                "fallback on the wider layer",
+            )
+
+    def test_matching_layer_agrees_with_generic_fallback(self):
+        """For a layer whose width == the global default, the hook and the
+        generic fallback must agree (no gratuitous divergence)."""
+        for module_name in ("qkv_proj", "o_proj"):
+            hook = self.model.get_hidden_dim(module_name, layer_idx=0)
+            generic = get_default_hidden_dim(module_name, self.model.config, 0)
+            self.assertEqual(hook, generic)
+
+    def test_missing_head_dim_raises_not_derives(self):
+        """Removing the fallback means an absent head_dim fails loudly instead
+        of silently deriving a wrong hidden_size // num_attention_heads."""
+        cfg = self.model.config
+        del cfg.head_dim
+        with self.assertRaises(AttributeError):
+            self.model.get_hidden_dim("o_proj", layer_idx=1)
+
+
+class TestLagunaNonAttentionDelegates(unittest.TestCase):
+    """Non-attention modules must delegate to the shared helper so that
+    `--lora-target-modules all` (MLP / MoE / embed / lm_head) still works
+    once the model defines `get_hidden_dim`.
+    """
+
+    def setUp(self):
+        self.model = _make_fake_laguna([48, 64])
+
+    def test_delegates_mlp_and_embedding_modules(self):
+        for module_name in (
+            "gate_up_proj",
+            "down_proj",
+            "gate_up_proj_moe",
+            "down_proj_moe",
+            "embed_tokens",
+            "lm_head",
+        ):
+            for layer_idx in (0, 1):
+                self.assertEqual(
+                    self.model.get_hidden_dim(module_name, layer_idx),
+                    get_default_hidden_dim(module_name, self.model.config, layer_idx),
+                    f"{module_name}@{layer_idx} should match the shared helper",
+                )
+
+    def test_unknown_module_raises(self):
+        with self.assertRaises(NotImplementedError):
+            self.model.get_hidden_dim("not_a_module", layer_idx=0)
+
+
+class TestLagunaForCausalLMDelegation(unittest.TestCase):
+    """`LagunaForCausalLM.get_hidden_dim` must forward to the inner model."""
+
+    def test_forwards_to_inner_model(self):
+        inner = _make_fake_laguna([48, 64])
+        causal = LagunaForCausalLM.__new__(LagunaForCausalLM)
+        # Bypass nn.Module.__setattr__: __new__ skips __init__, so the module
+        # registries it expects when assigning a Module-valued attr don't exist.
+        object.__setattr__(causal, "model", inner)
+        for module_name in ("qkv_proj", "o_proj", "gate_up_proj"):
+            for layer_idx in (0, 1):
+                self.assertEqual(
+                    causal.get_hidden_dim(module_name, layer_idx),
+                    inner.get_hidden_dim(module_name, layer_idx),
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()