From 754408f9f1664b68152421e57ae0753b99b1ceb2 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 11 Nov 2026 13:08:48 +0000 Subject: [PATCH] Use two row blocks for eight-row Kimi paired gathers A single thread block copies every row of the paired latent/router gather. Partitioning eight-row gathers across two blocks exposes more copy work while preserving each output's rank-major layout. Set row_blocks=2 only at eight physical rows and retain one block elsewhere. Pin the merged FlashInfer main that supplies the partitioned-gather API in both pyproject.toml and uv.lock. This ports the policy accepted in PR #258 onto main without replaying the counted-selection change already merged. Earlier isolated TP16 measurements reduced this gather from 33.31 to 31.00 microseconds; they do not establish full-model speedup. On this branch, 41 focused CPU checks pass, including five row-capacity policy cases. The dependency lock validates offline. GPU and end-to-end tests were not rerun for this port. Design-Signed-Off-By: Codex GPT-6; design-review=cfcf Signed-off-by: Martin Vit --- ignition/recipes/kimi3/routing.py | 2 +- pyproject.toml | 2 +- tests/recipes/test_kimi31.py | 26 ++++++++++++++++++++++++++ uv.lock | 4 ++-- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/ignition/recipes/kimi3/routing.py b/ignition/recipes/kimi3/routing.py index ffd96f7f..5ba9d861 100644 --- a/ignition/recipes/kimi3/routing.py +++ b/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", diff --git a/pyproject.toml b/pyproject.toml index 975792fa..1ef7c2f7 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/recipes/test_kimi31.py b/tests/recipes/test_kimi31.py index 6dec23d2..37a420a7 100644 --- a/tests/recipes/test_kimi31.py +++ b/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)): diff --git a/uv.lock b/uv.lock index df11638e..5216f094 100644 --- a/uv.lock +++ b/uv.lock @@ -400,7 +400,7 @@ source = { directory = "heliopause" } [[package]] name = "flashinfer-python" version = "0.7.3" -source = { git = "https://github.com/local-inference-lab/flashinfer.git?rev=485975daad1418358a2bf79340e3f785f0d48e04#485975daad1418358a2bf79340e3f785f0d48e04" } +source = { git = "https://github.com/local-inference-lab/flashinfer.git?rev=7bb37d608657e18a40de913d2cc3cd0448400336#7bb37d608657e18a40de913d2cc3cd0448400336" } dependencies = [ { name = "apache-tvm-ffi" }, { name = "click" }, @@ -563,7 +563,7 @@ requires-dist = [ { name = "cuda-python", marker = "extra == 'prepare'", specifier = "==13.3.1" }, { name = "cuda-tile", extras = ["tileiras"], marker = "extra == 'prepare'", specifier = "==1.6.0" }, { name = "flare-heliopause", marker = "extra == 'prepare'", directory = "heliopause" }, - { name = "flashinfer-python", extras = ["b12x"], marker = "extra == 'prepare'", git = "https://github.com/local-inference-lab/flashinfer.git?rev=485975daad1418358a2bf79340e3f785f0d48e04" }, + { name = "flashinfer-python", extras = ["b12x"], marker = "extra == 'prepare'", git = "https://github.com/local-inference-lab/flashinfer.git?rev=7bb37d608657e18a40de913d2cc3cd0448400336" }, { name = "huggingface-hub", marker = "extra == 'prepare'", specifier = "==1.33.0" }, { name = "openai", marker = "extra == 'api-test'", specifier = ">=2,<3" }, { name = "pillow", marker = "extra == 'prepare'", specifier = "==12.3.0" },