Demonstrate attention routing in Qwen 3

Author: DePasqualeOrg

Libraries/MLXLLM/Models/Qwen3.swift

@@ -76,17 +76,23 @@ private class Attention: Module {
         keys = kNorm(keys.reshaped(B, L, args.kvHeads, -1)).transposed(0, 2, 1, 3)
         values = values.reshaped(B, L, args.kvHeads, -1).transposed(0, 2, 1, 3)
 
+        // Apply RoPE positioning
         if let cache {
             queries = rope(queries, offset: cache.offset)
             keys = rope(keys, offset: cache.offset)
-            (keys, values) = cache.update(keys: keys, values: values)
         } else {
             queries = rope(queries)
             keys = rope(keys)
         }
 
-        let output = MLXFast.scaledDotProductAttention(
-            queries: queries, keys: keys, values: values, scale: scale, mask: mask
+        // Use the automatic attention router that handles both quantized and regular caches
+        let output = attentionWithCacheUpdate(
+            queries: queries,
+            keys: keys,
+            values: values,
+            cache: cache,
+            scale: scale,
+            mask: mask
         )
         .transposed(0, 2, 1, 3)
         .reshaped(B, L, -1)

mlx-swift-examples.xcodeproj/xcshareddata/xcschemes/llm-tool.xcscheme

@@ -61,8 +61,16 @@
          </CommandLineArgument>
          <CommandLineArgument
             argument = "--model mlx-community/Qwen2-VL-2B-Instruct-4bit --prompt &apos;Describe the image in English.&apos; --image https://www.gstatic.com/webp/gallery/1.webp"
+            isEnabled = "NO">
+         </CommandLineArgument>
+         <CommandLineArgument
+            argument = "--model mlx-community/Qwen3-1.7B-4bit --prompt &quot;Explain quantum computing in simple terms&quot; --max-tokens 100 --kv-bits 4"
             isEnabled = "YES">
          </CommandLineArgument>
+         <CommandLineArgument
+            argument = "--model mlx-community/Qwen3-1.7B-4bit --prompt &quot;Explain quantum computing in simple terms&quot; --max-tokens 100"
+            isEnabled = "NO">
+         </CommandLineArgument>
          <CommandLineArgument
             argument = "--repetition-penalty 1.2"
             isEnabled = "NO">