[sbryngelson]

Problem

aneforge/moe.py:_topk_gate builds the top-k mask by peeling the row maximum k times. Each round marked every expert tied at the maximum, so an exact fp16 probability tie activated more than k experts. On Qwen1.5-MoE-A2.7B (E=60, k=4) this happened in 36 of 1848 layer-steps of one 48-token decode, activating 5 or 6 experts.

Change

• Each round selects the lowest tied index (an amin over the tied positions) and removes only that expert, so exactly k experts are active per row.
• Expert indices are carried as fp16 integers, which are exact only through 2048. _topk_gate now rejects counts outside 1 <= k <= E <= 2049 with a clear error instead of silently selecting more than k.

Tests

• tests/test_moe.py::test_topk_gate_exact_ties_select_exactly_k (on-device): three tie patterns (a 3-way tie at the max, a tie at the k-th slot, a tie among the top two) and a no-tie row; checks the selected indices and that gate values are unchanged. On the previous code it fails (a 3-way tie with k=2 selects 6 experts).
• tests/test_moe_topk_range.py (off-device): rejected counts and the largest supported count.

Checks run (Apple M5 Pro, H17s, macOS 27.0.1)

• ruff check, the 2-space pylint gate, python -m compileall -q aneforge: clean.

• pyright aneforge: 0 errors with aneforge/_version.py present (generated at install). In this un-installed checkout it reports 2 errors for the missing generated file, identical on main.

• pytest -m "not requires_ane" -q: 525 passed, 1 skipped.

• tests/run_corpus.py: GREEN, 90/90.

• pytest tests/ -q (on-device, full): 27 failed, 1275 passed, 1 skipped. All 27 failures reproduce on unmodified main (4114158) on this machine: the same 27 tests fail on both trees when rerun serially, and the set did not change across runs. None involve MoE code. Most are ANE compile failures (ane_e5rt_program_compile failed; the native SDPA bridge reports verifyBundleAtPath: invalid model), which looks like a macOS 27.0.1 compatibility issue worth its own investigation:

• tests/test_autograd.py::test_transformer_block_trains

• tests/test_autograd.py::test_charlm_trains_and_predicts

• tests/test_autograd.py::test_charlm_generalizes_on_corpus

• tests/test_autograd.py::test_prenorm_transformer_block_trains

• tests/test_autograd.py::test_llama_block_trains

• tests/test_autograd.py::test_resident_transformer_block_trains

• tests/test_decoder_block.py::test_causal_llama_block_matches_numpy[6-16-2-32]

• tests/test_decoder_block.py::test_causal_llama_block_matches_numpy[8-32-4-64]

• tests/test_decoder_block.py::test_generate_fixed_cache_matches_per_step

• tests/test_linalg.py::test_lu_pivoted

• tests/test_linalg.py::test_det_matches_numpy[10.0]

• tests/test_linalg.py::test_det_matches_numpy[100.0]

• tests/test_linalg.py::test_slogdet_matches_numpy

• tests/test_linalg.py::test_solve_matches_numpy[4]

• tests/test_linalg.py::test_solve_matches_numpy[6]

• tests/test_linalg.py::test_solve_matches_numpy[8]

• tests/test_linalg.py::test_solve_matrix_rhs

• tests/test_linalg.py::test_solve_residual_is_small

• tests/test_linalg.py::test_inv_matches_numpy[4]

• tests/test_linalg.py::test_inv_matches_numpy[6]

• tests/test_linalg.py::test_inv_matches_numpy[8]

• tests/test_linalg.py::test_inv_times_a_is_identity

• tests/test_onnx.py::test_fuse_attention_causal_uses_native_sdpa

• tests/test_onnx.py::test_attention_matches_onnxruntime

• tests/test_sdpa_causal.py::test_af_sdpa_is_causal_native_end_to_end

• tests/test_sdpa_causal.py::test_native_sdpa_additive_mask_bridge

• tests/test_streaming.py::test_streamed_grads_match_monolith