[LoRA] Support DSA indexer LoRA targets for GLM-5.1 / DeepSeek-V3.2-family models (#28110)

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -1382,10 +1382,17 @@ class Indexer(MultiPlatformOp):
             topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
             return maybe_capture_indexer_topk(layer_id, topk_result)
 
+        # When weights_proj is LoRA-wrapped, use an eager module call so the
+        # wrapper owns base+delta and no LoRA kernel runs under torch.compile
+        weights_proj_lora = getattr(self.weights_proj, "set_lora", False)
+
         if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle():
             current_stream = torch.cuda.current_stream()
             self.alt_stream.wait_stream(current_stream)
-            weights = self._project_and_scale_head_gates(x)
+            if weights_proj_lora:
+                weights = self.weights_proj(x)[0].float() * self.n_heads**-0.5
+            else:
+                weights = self._project_and_scale_head_gates(x)
             query, key = self._get_q_k_bf16(
                 q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
             )
@@ -1472,6 +1479,11 @@ class Indexer(MultiPlatformOp):
                 x_for_gate = x
 
             if is_in_tc_piecewise_cuda_graph():
+                if weights_proj_lora:
+                    raise RuntimeError(
+                        "DSA indexer weights_proj LoRA is incompatible with TC piecewise CUDA graph; remove the explicit"
+                        " prefill cuda-graph backend override or drop indexer.weights_proj from the LoRA target modules."
+                    )
                 weights = logits_head_gate_pcg(
                     x_for_gate,
                     self.weights_proj.weight,
@@ -1479,6 +1491,9 @@ class Indexer(MultiPlatformOp):
                     self.softmax_scale,
                     q_scale,
                 )
+            elif weights_proj_lora:
+                weights = self.weights_proj(x_for_gate)[0].float() * self.n_heads**-0.5
+                weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
             else:
                 weights = self._get_logits_head_gate(x_for_gate, q_scale)
 

python/sglang/srt/lora/lora_manager.py

@@ -813,7 +813,11 @@ class LoRAManager:
                 continue
 
             # The module should be converted if it is included in target_names
-            if module_name.split(".")[-1] in self.target_modules:
+            parts = module_name.split(".")
+            if (
+                parts[-1] in self.target_modules
+                or ".".join(parts[-2:]) in self.target_modules
+            ):
                 layer_id = get_layer_id(module_name)
                 if layer_id is None:
                     continue

python/sglang/srt/lora/utils.py

@@ -185,6 +185,21 @@ def get_hidden_dim(
                 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)
@@ -248,6 +263,12 @@ def get_normalized_target_modules(
         "unembed_tokens": "lm_head",
         "q_a_proj": "fused_qkv_a_proj_with_mqa",
         "kv_a_proj_with_mqa": "fused_qkv_a_proj_with_mqa",
+        # DSA indexer projections are qualified with their parent module name
+        # because the bare leaf names collide with unrelated modules in other
+        # models (e.g. DeepSeek-V4 attention `wq_b`, Pixtral vision `wk`).
+        "wq_b": "indexer.wq_b",
+        "wk": "indexer.wk",
+        "weights_proj": "indexer.weights_proj",
     }
 
     result = set()
@@ -302,10 +323,14 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
 
 EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
 ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "out_proj", "down_proj", "down_proj_moe"]
+DSA_INDEXER_LORA_NAMES = frozenset(
+    {"indexer.wq_b", "indexer.wk", "indexer.weights_proj"}
+)
 REPLICATED_LINEAR_LORA_NAMES = [
     "fused_qkv_a_proj_with_mqa",
     "fc1_latent_proj",
     "fc2_latent_proj",
+    *DSA_INDEXER_LORA_NAMES,
 ]
 
 # Normalized module names that the LoRA system fully supports
@@ -328,6 +353,7 @@ _KNOWN_LORA_TARGET_MODULES = frozenset(
         "q_b_proj",
         "kv_b_proj",
     }
+    | DSA_INDEXER_LORA_NAMES
 )
 
 
@@ -347,6 +373,9 @@ def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set:
     )
 
     raw_names: set = set()
+    dsa_indexer_leaf_names = {
+        target_name.split(".")[-1] for target_name in DSA_INDEXER_LORA_NAMES
+    }
     for name, module in model.named_modules():
         if isinstance(module, FusedMoE):
             raw_names.add("gate_up_proj")
@@ -356,7 +385,18 @@ def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set:
         elif isinstance(module, VocabParallelEmbedding):
             raw_names.add("embed_tokens")
         elif isinstance(module, LinearBase):
-            raw_names.add(name.split(".")[-1])
+            parts = name.split(".")
+            leaf_name = parts[-1]
+            parent_qualified_name = ".".join(parts[-2:])
+            if parent_qualified_name in DSA_INDEXER_LORA_NAMES:
+                raw_names.add(parent_qualified_name)
+            elif leaf_name in dsa_indexer_leaf_names:
+                # Bare DSA indexer leaf names are ambiguous across model
+                # families. Only auto-detect them when the actual module path
+                # proves they are under an `indexer` parent.
+                continue
+            else:
+                raw_names.add(leaf_name)
 
     normalized = get_normalized_target_modules(raw_names)
     result = normalized & _KNOWN_LORA_TARGET_MODULES

python/sglang/srt/utils/common.py

@@ -3681,6 +3681,9 @@ SUPPORTED_LORA_TARGET_MODULES = [
     "kv_a_proj_with_mqa",
     "q_b_proj",
     "kv_b_proj",
+    "wq_b",
+    "wk",
+    "weights_proj",
     "gate_proj",
     "up_proj",
     "down_proj",