From 19d3f86895f43eb23b71a707726810cd700f947d Mon Sep 17 00:00:00 2001 From: Filip <34603115+PheelaV@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:48:05 +0100 Subject: [PATCH] [LoRA] Laguna: per-layer LoRA hidden-dim resolution for packed attention (#30298) Co-authored-by: Claude Opus 4.8 (1M context) --- python/sglang/srt/lora/utils.py | 200 +++++++++--------- python/sglang/srt/models/laguna.py | 29 +++ .../unit/lora/test_laguna_hidden_dim_unit.py | 163 ++++++++++++++ 3 files changed, 296 insertions(+), 96 deletions(-) create mode 100644 test/registered/unit/lora/test_laguna_hidden_dim_unit.py diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index c3a74a515..213e40a1a 100644 --- a/python/sglang/srt/lora/utils.py +++ b/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( diff --git a/python/sglang/srt/models/laguna.py b/python/sglang/srt/models/laguna.py index 3287d9776..c7c85e3ee 100644 --- a/python/sglang/srt/models/laguna.py +++ b/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"), diff --git a/test/registered/unit/lora/test_laguna_hidden_dim_unit.py b/test/registered/unit/lora/test_laguna_hidden_dim_unit.py new file mode 100644 index 000000000..9db1878fd --- /dev/null +++ b/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()