perf(v4): vectorize Engram doc-aware ngram gather (no Python batch loop)

cppmega_v4/nn/engram_v4.py::EngramV4Block._build_ngrams previously did
a Python 'for b_i in range(B)' loop to gather token_ids at clamped
positions per-batch when document_ids was provided. Replaced with one
mx.take_along_axis call on flattened positions.

For B=1 (typical decode) the change is a no-op; for larger batches
this collapses B kernel dispatches into one.

Verified: test_unified_superblock_v4.py — 9/9 (engram doc-ids tests
include the boundary-crossing assertion).

Author: apstenku123

cppmega_v4/nn/engram_v4.py

@@ -172,11 +172,13 @@ def _build_ngrams(
             # itself (the current token — guaranteed in-doc).
             current_pos = mx.broadcast_to(mx.arange(S)[None, :, None], (B, S, max_n))
             clamped_pos = mx.where(same_doc, positions[None, :, :], current_pos)
-            # Gather token_ids at clamped positions: [B, S, max_n]
-            ngrams_per_b = []
-            for b_i in range(B):
-                ngrams_per_b.append(mx.take(token_ids[b_i], clamped_pos[b_i], axis=0))
-            return mx.stack(ngrams_per_b, axis=0)
+            # Gather token_ids at clamped positions: [B, S, max_n].
+            # Vectorized over batch via take_along_axis: tok_b expanded
+            # to [B, S*max_n] for fancy-indexing along axis=1 with the
+            # flattened clamped_pos.
+            flat_pos = clamped_pos.reshape(B, S * max_n)
+            ngrams_flat = mx.take_along_axis(token_ids, flat_pos, axis=1)
+            return ngrams_flat.reshape(B, S, max_n)
         # No doc_ids: simple gather (with the >=0 clamp).
         ngrams = token_ids[:, positions]
         return ngrams