test

tests/ops/binned_copy.py

@@ -40,7 +40,7 @@ def make_stress_expert_capacity_tests():
     for seq_len, hidden_size, num_experts, top_k in [
         (4096, 768, 32, 4),
     ]:
-        for expert_capacity in [65535, 70000, 100000, 1000000]:
+        for expert_capacity in [65535, 70000, 90000]:
             tests.append((seq_len, hidden_size, num_experts, top_k, expert_capacity))
 
     return tuple(tests)
@@ -80,7 +80,6 @@ def test_binned_gather(seq_len: int, hidden_size: int, num_experts: int, top_k:
                 # indices[start:end] are the indices for this expert
                 # For each slot j, get the input index and copy the row
                 idx = indices[start : start + num_tokens] // top_k
-                print(f"Expert {i}: indices[{start}:{start + num_tokens}] = {indices[start : start + num_tokens]} -> tokens {idx}")
                 out[i, :num_tokens, :] = x[idx, :]
             start = end
         return out