Merge pull request #1505 from kvcache-ai/support-qwen3next

fix qwen3next bug

Author: ovowei

doc/en/Qwen3-Next.md

@@ -38,6 +38,7 @@ To install KTransformers, follow the official [Installation Guide](https://kvcac
 python ktransformers/server/main.py \
   --port 10021 \
   --model_path path-to-Qwen3-Next-80B-A3B-Thinking \
+  --gguf_path path-to-Qwen3-Next-80B-A3B-Thinking \
   --model_name Qwen3NextForCausalLM \
   --optimize_config_path <local_path>/ktransformers/optimize/optimize_rules/Qwen3Next-serve.yaml \
   --max_new_tokens 1024 \

ktransformers/models/modeling_qwen3_next.py

@@ -43,10 +43,23 @@
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from transformers.utils.deprecation import deprecate_kwarg
 from transformers.utils.generic import OutputRecorder, check_model_inputs
-from transformers.utils.import_utils import (
-    is_causal_conv1d_available,
-    is_flash_linear_attention_available,
-)
+try:
+    from transformers.utils.import_utils import (
+        is_causal_conv1d_available,
+        is_flash_linear_attention_available,
+    )
+except ImportError:
+    is_causal_conv1d_available = lambda: False
+
+
+try:
+    from transformers.utils.import_utils import (
+        is_flash_linear_attention_available,
+    )
+except ImportError:
+    is_flash_linear_attention_available = lambda: False
+
+
 from .configuration_qwen3_next import Qwen3NextConfig
 
 

ktransformers/operators/balance_serve_attention.py

@@ -614,11 +614,6 @@ def forward(self,
             query_states = self.q_norm(query_states, bsz_tensors)
             key_states = self.k_norm(key_states, bsz_tensors)
 
-
-        query_states = query_states.view(q_len, self.config.num_attention_heads, self.head_dim)
-        key_states = key_states.view(q_len, self.config.num_key_value_heads, self.head_dim)
-        value_states = value_states.view(q_len, self.config.num_key_value_heads, self.head_dim)
-        
         # cos, sin = freqs_cis
         """
         print(query_states.shape)
@@ -634,11 +629,16 @@ def forward(self,
         if freqs_cis is not None:  
             query_states, key_states = self.apply_rotary_pos_emb(query_states.unsqueeze(0), key_states.unsqueeze(0), freqs_cis)
 
+        query_states = query_states.view(q_len, self.config.num_attention_heads, self.head_dim)
+        key_states = key_states.view(q_len, self.config.num_key_value_heads, self.head_dim)
+        value_states = value_states.view(q_len, self.config.num_key_value_heads, self.head_dim)
+
 
         k_cache = kv_cache.get_k_cache(self.layer_idx)
         v_cache = kv_cache.get_v_cache(self.layer_idx)
 
 
+        print(f"{k_cache.shape=}, {v_cache.shape=}, {query_states.shape=}, {key_states.shape=}, {value_states.shape=}")
         attn_output = wrapper.forward(query_states, k_cache, v_cache, key_states, value_states)
 
 

ktransformers/server/args.py

@@ -3,6 +3,7 @@
 from ktransformers.util.utils import get_free_ports
 from transformers import AutoConfig
 from ktransformers.models.configuration_qwen3_moe import Qwen3MoeConfig
+from ktransformers.models.configuration_qwen3_next import Qwen3NextConfig
 from ktransformers.models.configuration_smallthinker import SmallthinkerConfig
 from ktransformers.models.configuration_glm4_moe import Glm4MoeConfig
 
@@ -138,12 +139,16 @@ def parse_args(self):
         self.cfg.server_ip = args.host
         self.cfg.server_port = args.port
         self.cfg.user_force_think = args.force_think
+
+
+        args.architectures = args.model_name
+
         try:
             model_config = AutoConfig.from_pretrained(args.model_dir, trust_remote_code=True)
         except:
-            try:
-                model_config = Glm4MoeConfig.from_pretrained(args.model_dir, trust_remote_code=True)
-            except:
+            if args.model_name == "Qwen3NextForCausalLM":
+                model_config = Qwen3NextConfig.from_pretrained(args.model_dir)
+            else:
                 raise ValueError(f"Model {args.model_name} not supported. Please check your model directory or model name.")