[II] Preserve FP32 and FP8 Kimi projection payloads

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)

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: