diff --git a/tests/models/kimi_k3/test_tp_projection.py b/tests/models/kimi_k3/test_tp_projection.py index bd733076effa..150beb3c5bc1 100644 --- a/tests/models/kimi_k3/test_tp_projection.py +++ b/tests/models/kimi_k3/test_tp_projection.py @@ -102,6 +102,74 @@ def gather(transport, projection_group, *, max_batch_size): assert received["transport"].shape == (1, 1, 136) +@pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.") +def test_b12x_projection_gather_preserves_fp32_payload_bits(monkeypatch): + tp_size = 2 + local = torch.tensor([[1.25, -2.5]], dtype=torch.float32, device="cuda") + other = torch.tensor([[3.75, -4.5]], dtype=torch.float32, device="cuda") + group = SimpleNamespace(world_size=tp_size, ranks=list(range(tp_size))) + received: dict[str, object] = {} + monkeypatch.setattr(tp_projection.envs, "VLLM_USE_B12X_DCP_A2A", True) + monkeypatch.setattr( + tp_projection, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr(tp_projection, "_get_kimi_projection_group", lambda: group) + + def gather(transport, projection_group, *, max_batch_size): + received["transport"] = transport + result = torch.empty( + (1, tp_size, transport.shape[-1]), + dtype=transport.dtype, + device=transport.device, + ) + result[0, 0].copy_(transport[0, 0]) + result[0, 1].copy_(other.view(torch.float8_e4m3fn).flatten()) + return result + + monkeypatch.setattr(tp_projection, "dcp_b12x_all_gather_heads", gather) + + actual = tp_projection.gather_kimi_sharded_projection(local) + + torch.testing.assert_close(actual, torch.cat((local, other), dim=-1)) + assert actual.dtype == torch.float32 + assert received["transport"].dtype == torch.float8_e4m3fn + assert received["transport"].shape == (1, 1, 8) + + +@pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.") +def test_b12x_projection_gather_preserves_fp8_payload(monkeypatch): + tp_size = 2 + local = ( + torch.arange(16, dtype=torch.float32, device="cuda") + .to(torch.float8_e4m3fn) + .view(1, -1) + ) + other = ( + torch.arange(16, 32, dtype=torch.float32, device="cuda") + .to(torch.float8_e4m3fn) + .view(1, -1) + ) + group = SimpleNamespace(world_size=tp_size, ranks=list(range(tp_size))) + monkeypatch.setattr(tp_projection.envs, "VLLM_USE_B12X_DCP_A2A", True) + monkeypatch.setattr( + tp_projection, "get_tensor_model_parallel_world_size", lambda: tp_size + ) + monkeypatch.setattr(tp_projection, "_get_kimi_projection_group", lambda: group) + + def gather(transport, projection_group, *, max_batch_size): + return torch.stack((transport[:, 0], other), dim=1) + + monkeypatch.setattr(tp_projection, "dcp_b12x_all_gather_heads", gather) + + actual = tp_projection.gather_kimi_sharded_projection(local) + + torch.testing.assert_close( + actual.float(), + torch.cat((local, other), dim=-1).float(), + ) + assert actual.dtype == torch.float8_e4m3fn + + def test_projection_gather_uses_standard_collective_outside_decode(monkeypatch): local = torch.arange(12).view(3, 4) expected = torch.cat((local, local), dim=-1) diff --git a/vllm/models/kimi_k3/nvidia/tp_projection.py b/vllm/models/kimi_k3/nvidia/tp_projection.py index 2a8638290e20..5f90449d1894 100644 --- a/vllm/models/kimi_k3/nvidia/tp_projection.py +++ b/vllm/models/kimi_k3/nvidia/tp_projection.py @@ -43,7 +43,6 @@ def _try_b12x_kimi_projection_gather( not envs.VLLM_USE_B12X_DCP_A2A or output_parallel.ndim != 2 or output_parallel.shape[0] != 1 - or output_parallel.dtype not in (torch.float16, torch.bfloat16) or not output_parallel.is_cuda or not output_parallel.is_contiguous() ): @@ -55,15 +54,31 @@ def _try_b12x_kimi_projection_gather( projection_group = _get_kimi_projection_group() local_width = output_parallel.shape[1] + restore_dtype: torch.dtype | None = None strip_local_width: int | None = None - if local_width % 8 == 0: + if output_parallel.dtype in (torch.float16, torch.bfloat16): + if local_width % 8 == 0: + transport = output_parallel.view(1, 1, local_width) + else: + padded_width = (local_width + 7) // 8 * 8 + transport = torch.nn.functional.pad( + output_parallel, (0, padded_width - local_width) + ).view(1, 1, padded_width) + strip_local_width = local_width + elif output_parallel.dtype == torch.float32: + raw_width = local_width * output_parallel.element_size() + if raw_width % 8 != 0: + return None + # The FP8 view exposes one-byte transport lanes without converting the + # FP32 payload. The gathered result is restored to the original dtype. + transport = output_parallel.view(torch.float8_e4m3fn).view(1, 1, raw_width) + restore_dtype = torch.float32 + elif output_parallel.dtype == torch.float8_e4m3fn: + if local_width % 16 != 0: + return None transport = output_parallel.view(1, 1, local_width) else: - padded_width = (local_width + 7) // 8 * 8 - transport = torch.nn.functional.pad( - output_parallel, (0, padded_width - local_width) - ).view(1, 1, padded_width) - strip_local_width = local_width + return None gathered = dcp_b12x_all_gather_heads( transport, @@ -72,7 +87,10 @@ def _try_b12x_kimi_projection_gather( ) if strip_local_width is not None: gathered = gathered.narrow(-1, 0, strip_local_width).contiguous() - return gathered.flatten(1) + gathered = gathered.flatten(1) + if restore_dtype is not None: + gathered = gathered.view(restore_dtype) + return gathered def gather_kimi_sharded_projection(output_parallel: torch.Tensor) -> torch.Tensor: