diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 1ca901460..0f5b3c050 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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) diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index a43c5dfce..91a9e48d5 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/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 diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index b2984ed7d..e8301d5b4 100644 --- a/python/sglang/srt/lora/utils.py +++ b/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 diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 9c349b04b..ac44ba8ee 100644 --- a/python/sglang/srt/utils/common.py +++ b/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",