fix: pipeline sharding corruption for hybrid SSM/attention models

[adamteale]

Hybrid models (Qwen3.5/3.6/3.8, Kimi-Linear, Nemotron) with interleaved linear-attention (SSM) and full-attention layers were silently corrupted under 2-node pipeline sharding. Five root causes identified and fixed:

1. PipelineLastLayer: prefill-final single-token step now all_gathers so all ranks sample the same first decode token. Without this, ranks could argmax different tokens at >64 prompt tokens, diverging from step 0.

2. pipeline_parallel_prefill: process final prompt token exactly ONCE (was twice), preventing SSM recurrent-state double-advance that permanently corrupted hybrid models.

3. prefill() snapshot restore: use snapshots[-1] (last per-chunk = state after prompt[:-2]) instead of snapshots[-2], which for single-chunk prefills pointed at a mid-prompt state and rolled SSM back incorrectly.

4. batch_generate: wrap mlx_gen.insert() (the [-2:] decode restart) with set_pipeline_prefill(True/False) so the internal 2-token prefill uses the same no-all_gather path as the original prefill. all_gather during restart produced different bf16 reduction orders, corrupting long prompts.

5. Added test_pipeline_bit_exact.py: synthetic 4-layer qwen3_5_moe model with bit-exact comparison between single-process and 2-rank pipeline across prompt lengths 64-2048. All sizes now pass.

Verified on live 2-Mac TB5 cluster: Qwen3.6-35B-A3B-4bit 2-node sharded produces spec-exact code at all prompt lengths (was 0/9 garbled before). Decode throughput: ~71 tok/s sharded (faster than single-node 60 tok/s).

Motivation

Changes

Why It Works

Test Plan

Manual Testing

Automated Testing