Improve linear.py to load sharded weights & remove the dependency of Parameters from vllm (#2784)

Co-authored-by: SangBin Cho rkooo567@gmail.com

python/sglang/srt/layers/attention/__init__.py

@@ -66,7 +66,14 @@ class AttentionBackend(ABC):
         if forward_batch.forward_mode.is_decode():
             return self.forward_decode(q, k, v, layer, forward_batch, save_kv_cache)
         else:
-            return self.forward_extend(q, k, v, layer, forward_batch, save_kv_cache)
+            return self.forward_extend(
+                q,
+                k,
+                v,
+                layer,
+                forward_batch,
+                save_kv_cache,
+            )
 
     def forward_decode(
         self,

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -347,6 +347,8 @@ class FlashInferAttnBackend(AttentionBackend):
             else forward_batch.encoder_out_cache_loc
         )
 
+        logits_soft_cap = layer.logit_cap
+
         if not self.forward_metadata.use_ragged:
             if k is not None:
                 assert v is not None
@@ -359,7 +361,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 causal=not layer.is_cross_attention,
                 sm_scale=layer.scaling,
                 window_left=layer.sliding_window_size,
-                logits_soft_cap=layer.logit_cap,
+                logits_soft_cap=logits_soft_cap,
             )
         else:
             o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
@@ -368,7 +370,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 v.contiguous().view(-1, layer.tp_v_head_num, layer.head_dim),
                 causal=True,
                 sm_scale=layer.scaling,
-                logits_soft_cap=layer.logit_cap,
+                logits_soft_cap=logits_soft_cap,
             )
 
             if self.forward_metadata.extend_no_prefix: