Use two row blocks for eight-row Kimi paired gathers

ignition/recipes/kimi3/routing.py

@@ -72,7 +72,7 @@ def __init__(self, local_down, local_router, correction_bias, latent, weights, i
         else:
             self.logits = torch.zeros((rows, 896), dtype=torch.float32, device=latent.device)
             self.pair = dict(local_first=local_down, local_second=local_router, out_first=latent,
-                             out_second=self.logits, threads=512)
+                             out_second=self.logits, threads=512, row_blocks=2 if rows == 8 else 1)
             self.select = dict(router_logits=self.logits, correction_bias=correction_bias,
                                output_weights=weights, output_ids=ids)
             self.plan = pcie.plan(pcie.query_from_runtime(runtime, surface="DcpAllToAll.all_gather_pair",

pyproject.toml

@@ -22,7 +22,7 @@ prepare = [
     "torch==2.14.1+cu132",
     "torchvision==0.29.1+cu132",
     "transformers[torch,vision,chat_template] @ git+https://github.com/huggingface/transformers.git@1d9654e78a23fb167ef539cbd699dfbd3e4830cf",
-    "flashinfer-python[b12x] @ git+https://github.com/local-inference-lab/flashinfer.git@485975daad1418358a2bf79340e3f785f0d48e04",
+    "flashinfer-python[b12x] @ git+https://github.com/local-inference-lab/flashinfer.git@7bb37d608657e18a40de913d2cc3cd0448400336",
     "cuda-python==13.3.1",
     "cuda-tile[tileiras]==1.6.0",
     "safetensors==0.8.0",

tests/recipes/test_kimi31.py

@@ -209,6 +209,32 @@ def test_rank_preparation_rejects_invalid_k31_inputs_before_cuda(raw_config, tmp
         prepare_rank(tmp_path, PreparationConfig(**options), rank=0, group=None, core_storage=None)
 
 
+@pytest.mark.parametrize("rows,expected_blocks", [(2, 1), (4, 1), (6, 1), (7, 1), (8, 2)])
+def test_paired_route_gather_uses_two_row_blocks_only_for_p8(monkeypatch, rows, expected_blocks):
+    from types import SimpleNamespace
+    import torch
+    from b12x.comm import pcie
+    from ignition.recipes.kimi3 import routing
+
+    class Plan:
+        def request(self, *args, **kwargs):
+            return object()
+
+    monkeypatch.setattr(torch.distributed, "get_process_group_ranks", lambda group: tuple(range(16)))
+    monkeypatch.setattr(pcie, "query_from_runtime",
+                        lambda runtime, *, surface, call: dict(surface=surface, call=call))
+    monkeypatch.setattr(pcie, "plan", lambda query, *, runtime: Plan())
+    monkeypatch.setattr(routing, "declare", lambda *args, **kwargs: None)
+    runtime = SimpleNamespace(world_size=16, exchange_group=object())
+    local_down = torch.zeros((rows, 224), dtype=torch.bfloat16)
+    local_router = torch.zeros((rows, 56), dtype=torch.float32)
+    operation = routing.PreparedRouteGather(
+        local_down, local_router, torch.zeros(896), torch.zeros((rows, 3584), dtype=torch.bfloat16),
+        torch.zeros((rows, 16)), torch.zeros((rows, 16), dtype=torch.int32),
+        runtime=runtime, session=None, coordinator=None, name="route")
+    assert operation.pair["row_blocks"] == expected_blocks
+
+
 def test_native_quantization_covers_all_routed_experts(config):
     assert config.text["quantization_config"].get("weight_scale_encoding") is None
     for layer in range(config.text["first_k_dense_replace"], len(config.layers)):