[MLX] Qwen return last-token logits from forward; make SamplingHead operate on (B, vocab) (#20604)

Summary

Qwen's MLX forward now returns last-token logits (B, vocab) instead of
(B, S, vocab), so lm_head runs on one position per prefill instead of
the whole sequence. SamplingHead correspondingly drops its internal
last-token slice and samples directly from (B, vocab), making it
model-agnostic.

Changes

- backends/mlx/llm/sampling.py — SamplingHead drops its internal
logits[:, -1, :] slice; samples a token (B) directly from (B, vocab),
removing the (B, S, vocab) path.

- examples/models/qwen3_5_moe/export.py — _clean_forward returns
lm_head(x[:, -1, :]) → (B, vocab) instead of all S positions, so lm_head
runs once per prefill (chunk) instead of over the whole sequence.

- run.py + test_chunked_prefill.py — consumers of qwen's forward output,
both dropping the [0, -1, :] indexing.

- test_ops.py + test_sample.py — op-level tests of SamplingHead/sample
shape + node counts (upstream's top-k PR).

- No C++ changes — the runtime sampler already handles 2D; chunked
prefill result is identical.
  
Testing
 - MLX sample op tests pass on-device.
- Prefill ~19% faster (1505-token prompt): ~1352 → ~1612 tok/s; decode
unchanged (~75 tok/s).

Author: kiymetakdemir

backends/mlx/llm/sampling.py

@@ -12,7 +12,8 @@
 
 class SamplingHead(nn.Module):
     """
-    Wraps a model that returns logits and samples a token id on-device.
+    Wraps a model that returns last-token logits ``(B, vocab)`` and samples a
+    token id ``(B)`` on-device.
 
         forward(*model_args, temperature, top_k=None, top_p=1.0, seed=None,
                 **model_kwargs) -> token_id
@@ -33,12 +34,11 @@ def __init__(self, model: nn.Module):
         self.model = model
 
     def forward(self, *args, temperature, top_k=None, top_p=1.0, seed=None, **kwargs):
-        logits = self.model(*args, **kwargs)  # [B, S, vocab]
-        last = logits[:, -1, :]  # [B, vocab]
+        logits = self.model(*args, **kwargs)  # [B, vocab]
         if not isinstance(top_p, torch.Tensor):
             top_p = torch.tensor(float(top_p))
         if top_k is None:
             top_k = torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64)
         elif not isinstance(top_k, torch.Tensor):
             top_k = torch.tensor(int(top_k), dtype=torch.int64)
-        return torch.ops.mlx.sample(last, temperature, top_k, top_p, seed)
+        return torch.ops.mlx.sample(logits, temperature, top_k, top_p, seed)

backends/mlx/test/test_ops.py

@@ -7694,7 +7694,7 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
 
 
 class _LogitsPassthrough(nn.Module):
-    """Stand-in for a model returning logits [B, S, vocab]."""
+    """Stand-in for a model returning logits [B, vocab]."""
 
     def forward(self, logits: torch.Tensor) -> torch.Tensor:
         return logits
@@ -7759,7 +7759,7 @@ class SampleSeededTest(OpTestCase):
         "SortNode": 2,  # top-k threshold + top-p nucleus chain
         "CumsumNode": 1,
         "MinNode": 1,
-        "TakeNode": 2,  # last-token slice + top-k threshold gather
+        "TakeNode": 1,  # top-k threshold gather
         "ExpandDimsNode": 1,
         "WhereNode": 3,
     }
@@ -7769,7 +7769,7 @@ def create_model(self) -> nn.Module:
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (
-            torch.randn(1, 4, 256),
+            torch.randn(1, 256),
             torch.tensor(0.8),
             torch.tensor(0, dtype=torch.int64),
         )
@@ -7793,7 +7793,7 @@ def create_model(self) -> nn.Module:
         return UnseededSampleModel()
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
-        return (torch.randn(1, 4, 256), torch.tensor(0.8))
+        return (torch.randn(1, 256), torch.tensor(0.8))
 
 
 @register_test
@@ -7811,7 +7811,7 @@ class SampleTopPTest(OpTestCase):
         "SortNode": 2,
         "CumsumNode": 1,
         "MinNode": 1,
-        "TakeNode": 2,  # last-token slice + top-k threshold gather
+        "TakeNode": 1,  # top-k threshold gather
         "ExpandDimsNode": 1,
         "WhereNode": 3,
     }
@@ -7821,7 +7821,7 @@ def create_model(self) -> nn.Module:
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (
-            torch.randn(1, 4, 256),
+            torch.randn(1, 256),
             torch.tensor(0.8),
             torch.tensor(0, dtype=torch.int64),
             torch.tensor(0.9),
@@ -7843,7 +7843,7 @@ class SampleTopKTest(OpTestCase):
         "SortNode": 2,
         "CumsumNode": 1,
         "MinNode": 1,
-        "TakeNode": 2,  # last-token slice + top-k threshold gather
+        "TakeNode": 1,  # top-k threshold gather
         "ExpandDimsNode": 1,
         "LogicalOrNode": 0,
         "WhereNode": 3,
@@ -7854,7 +7854,7 @@ def create_model(self) -> nn.Module:
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         return (
-            torch.randn(1, 4, 256),
+            torch.randn(1, 256),
             torch.tensor(0.8),
             torch.tensor(0, dtype=torch.int64),
             torch.tensor(2, dtype=torch.int64),
@@ -7895,9 +7895,9 @@ def create_model(self) -> nn.Module:
         return SeededSampleModel()
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
-        logits = torch.randn(self.batch, 4, 1024, dtype=self.dtype)
+        logits = torch.randn(self.batch, 1024, dtype=self.dtype)
         if self.dtype == torch.bfloat16:
-            logits[0, -1, 512] = 50.0  # dominant -> unambiguous bf16 argmax
+            logits[0, 512] = 50.0  # dominant -> unambiguous bf16 argmax
         return (
             logits,
             torch.tensor(self.temperature),

backends/mlx/test/test_sample.py

@@ -35,7 +35,7 @@
 
 
 class _LogitsPassthrough(nn.Module):
-    """Stand-in for a model returning logits [B, S, vocab]."""
+    """Stand-in for a model returning logits [B, vocab]."""
 
     def forward(self, logits: torch.Tensor) -> torch.Tensor:
         return logits
@@ -205,7 +205,7 @@ def test_runtime_temperature_single_export(self):
         batch = 256
         torch.manual_seed(0)
         row = torch.randn(vocab)
-        logits = row.expand(batch, 1, vocab).contiguous()  # [B, S=1, vocab]
+        logits = row.expand(batch, vocab).contiguous()  # [B, vocab]
         seed = torch.tensor(0, dtype=torch.int64)
 
         run = torch.export.export(
@@ -224,7 +224,7 @@ def test_seeded_export_reproducible_no_host_rng(self):
         # exported program, independent of host RNG state (the seed is a graph
         # input, not host-side stateful RNG). Different seed -> different draws.
         torch.manual_seed(0)
-        logits = torch.randn(128, 1, 64)
+        logits = torch.randn(128, 64)
         seed = torch.tensor(123, dtype=torch.int64)
 
         run = torch.export.export(
@@ -250,7 +250,7 @@ def setUp(self):
 
     def test_top_p_end_to_end(self):
         # On-device nucleus: probs [0.5,0.3,0.15,0.05], top_p=0.9 -> token in {0,1,2}.
-        logits = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])).view(1, 1, 4)
+        logits = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])).view(1, 4)
         inputs = (
             logits,
             torch.tensor(1.0),
@@ -268,7 +268,7 @@ def test_top_p_end_to_end(self):
 
     def test_top_k_end_to_end(self):
         # On-device top-k: probs [0.5,0.3,0.15,0.05], top_k=2 -> token in {0,1}.
-        logits = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])).view(1, 1, 4)
+        logits = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])).view(1, 4)
         inputs = (
             logits,
             torch.tensor(1.0),

examples/models/qwen3_5_moe/export.py

@@ -728,7 +728,7 @@ def _clean_forward(self, tokens, input_pos):
         for layer in self.layers:
             x = layer(x, input_pos)
         x = self.norm(x)
-        return self.lm_head(x)
+        return self.lm_head(x[:, -1, :])
 
     model.forward = types.MethodType(_clean_forward, model)
 

examples/models/qwen3_5_moe/run.py

@@ -51,7 +51,7 @@ def _next_token(outputs, use_sampling: bool, temperature: float) -> int:
     """A --sample model returns the token id directly; else sample from logits."""
     if use_sampling:
         return int(outputs[0].reshape(-1)[0].item())
-    logits = outputs[0][0, -1, :]
+    logits = outputs[0][0]
     if temperature > 0:
         return int(torch.multinomial(torch.softmax(logits / temperature, dim=-1), 1))
     return int(torch.argmax(logits))

examples/models/qwen3_5_moe/test_chunked_prefill.py

@@ -52,8 +52,8 @@ def _scalar_metadata(program, name, default):
 
 
 def _last_logits(outputs):
-    # forward returns logits shaped (1, T, vocab); take the final position.
-    return outputs[0][0, -1, :]
+    # forward returns last-token logits shaped (1, vocab).
+    return outputs[0][0]
 
 
 class TestChunkedPrefill(unittest.TestCase):