Add TP3 lowbit and profiling probes

Author: s0l

vllm/distributed/device_communicators/lowbit_all_reduce.py

@@ -21,6 +21,9 @@
 
 logger = init_logger(__name__)
 _INT8_WORKSPACES: dict[tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor]] = {}
+_INT8_BF16PACK_WORKSPACES: dict[
+    tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor, int]
+] = {}
 _FP4_WORKSPACES: dict[tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor]] = {}
 _SHARED_ROW_WORKSPACES: dict[tuple[int, int, int], tuple[torch.Tensor, torch.Tensor]] = {}
 
@@ -287,6 +290,7 @@ def should_use_tp3_lowbit_hidden_reduce(
             "fp4_block16",
             "fp4_block16_packed",
             "int8_shared_row",
+            "int8_bf16pack",
         )
         and should_use_tp3_int8_hidden_reduce(input_, world_size, min_tokens)
     )
@@ -328,6 +332,29 @@ def _int8_workspace(
     return send, gathered
 
 
+def _int8_bf16pack_workspace(
+    device: torch.device,
+    world_size: int,
+    rows: int,
+    cols: int,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    device_index = device.index if device.index is not None else torch.cuda.current_device()
+    q_nbytes = rows * cols
+    s_nbytes = rows * 4
+    total_nbytes = q_nbytes + s_nbytes
+    packed_numel = (total_nbytes + 1) // 2
+    key = (device_index, world_size, rows, cols)
+    cached = _INT8_BF16PACK_WORKSPACES.get(key)
+    if cached is not None and cached[0].numel() == packed_numel:
+        return cached
+    send = torch.empty(packed_numel, device=device, dtype=torch.bfloat16)
+    gathered = torch.empty((world_size, packed_numel), device=device,
+                           dtype=torch.bfloat16)
+    cached = (send, gathered, total_nbytes)
+    _INT8_BF16PACK_WORKSPACES[key] = cached
+    return cached
+
+
 def _fp4_workspace(
     device: torch.device,
     world_size: int,
@@ -403,6 +430,58 @@ def tp3_int8_hidden_all_reduce(
     return out
 
 
+def tp3_int8_bf16pack_hidden_all_reduce(
+    input_: torch.Tensor,
+    group: ProcessGroup,
+) -> torch.Tensor:
+    """Approximate SUM all-reduce with int8 bytes transported as BF16.
+
+    This does not use BF16 arithmetic. The BF16 tensor is only a 16-bit NCCL
+    transport container for the existing int8 codes plus fp32 row scales.
+    """
+    assert triton is not None
+    world_size = dist.get_world_size(group)
+    assert world_size == 3
+    rows, cols = input_.shape
+    trace, trace_sync = _trace_enabled(input_)
+    t0 = _mark_trace(trace, trace_sync, input_.device)
+    q_nbytes = rows * cols
+    send_bf16, gathered_bf16, total_nbytes = _int8_bf16pack_workspace(
+        input_.device, world_size, rows, cols
+    )
+    send = send_bf16.view(torch.uint8)[:total_nbytes]
+    gathered = gathered_bf16.view(torch.uint8).view(world_size, -1)[:, :total_nbytes]
+    q_local = send[:q_nbytes].view(rows, cols)
+    s_local = send[q_nbytes:].view(torch.float32)
+    block = triton.next_power_of_2(cols)
+    _quantize_i8_row_kernel[(rows,)](input_, q_local, s_local, cols, block)
+    t1 = _mark_trace(trace, trace_sync, input_.device)
+    dist.all_gather_into_tensor(gathered_bf16, send_bf16, group=group)
+    t2 = _mark_trace(trace, trace_sync, input_.device)
+    out = torch.empty_like(input_)
+    q_gathered = gathered[:, :q_nbytes].view(world_size, rows, cols)
+    s_gathered = gathered[:, q_nbytes:].view(torch.float32).view(world_size, rows)
+    _dequant_sum_i8_row_kernel[(rows,)](
+        q_gathered[0], q_gathered[1], q_gathered[2],
+        s_gathered[0], s_gathered[1], s_gathered[2],
+        out, cols, block,
+    )
+    t3 = _mark_trace(trace, trace_sync, input_.device)
+    if trace:
+        logger.warning(
+            "AG2_LOWBIT_REDUCE_TRACE mode=int8_bf16pack tokens=%d hidden=%d "
+            "quant=%.6f gather=%.6f dequant=%.6f total=%.6f payload_mib=%.3f",
+            rows,
+            cols,
+            t1 - t0,
+            t2 - t1,
+            t3 - t2,
+            t3 - t0,
+            send_bf16.numel() * send_bf16.element_size() / 2**20,
+        )
+    return out
+
+
 def tp3_int8_tensor_hidden_all_reduce(
     input_: torch.Tensor,
     group: ProcessGroup,
@@ -588,6 +667,8 @@ def tp3_lowbit_hidden_all_reduce(
     group: ProcessGroup,
     mode: str,
 ) -> torch.Tensor:
+    if mode == "int8_bf16pack":
+        return tp3_int8_bf16pack_hidden_all_reduce(input_, group)
     if mode == "int8_shared_row":
         return tp3_int8_shared_row_hidden_all_reduce(input_, group)
     if mode == "int8_tensor":

vllm/model_executor/layers/linear.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import itertools
+import os
+import time
 from abc import abstractmethod
 from collections.abc import Iterable
 from typing import Any
@@ -45,6 +47,33 @@
 
 logger = init_logger(__name__)
 
+
+def _ag2_row_linear_trace_bool(name: str, default: str = "0") -> bool:
+    return os.environ.get(name, default).strip().lower() in ("1", "true", "yes", "on")
+
+
+def _ag2_row_linear_trace_enabled(prefix: str, input_: torch.Tensor) -> bool:
+    if not _ag2_row_linear_trace_bool("AG2_VLLM_ROW_LINEAR_TRACE"):
+        return False
+    min_tokens = int(os.environ.get("AG2_VLLM_ROW_LINEAR_TRACE_MIN_TOKENS", "8192"))
+    if input_.dim() == 0 or int(input_.shape[0]) < min_tokens:
+        return False
+    prefix_filter = os.environ.get("AG2_VLLM_ROW_LINEAR_TRACE_PREFIX", "")
+    if prefix_filter and prefix_filter not in prefix:
+        return False
+    if _ag2_row_linear_trace_bool("AG2_VLLM_ROW_LINEAR_TRACE_RANK0_ONLY", "1"):
+        if os.environ.get("LOCAL_RANK") is not None:
+            return os.environ.get("LOCAL_RANK") == "0"
+        if torch.cuda.is_available():
+            return torch.cuda.current_device() == 0
+    return True
+
+
+def _ag2_row_linear_mark(enabled: bool, device: torch.device) -> float:
+    if enabled and _ag2_row_linear_trace_bool("AG2_VLLM_ROW_LINEAR_TRACE_SYNC"):
+        torch.cuda.synchronize(device)
+    return time.perf_counter()
+
 WEIGHT_LOADER_V2_SUPPORTED = [
     "UnquantizedLinearMethod",
     "CompressedTensorsLinearMethod",
@@ -2203,24 +2232,47 @@ def forward(
         self,
         input_,
     ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        trace_enabled = _ag2_row_linear_trace_enabled(self.prefix, input_)
+        device = input_.device
+        t0 = _ag2_row_linear_mark(trace_enabled, device) if trace_enabled else 0.0
         if self.input_is_parallel:
             input_parallel = input_
         else:
             split_input = split_tensor_along_last_dim(
                 input_, num_partitions=self.tp_size
             )
             input_parallel = split_input[self.tp_rank].contiguous()
+        t1 = _ag2_row_linear_mark(trace_enabled, device) if trace_enabled else 0.0
 
         # Matrix multiply.
         # Only fuse bias add into GEMM for rank 0 (this ensures that
         # bias will not get added more than once in TP>1 case)
         bias_ = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
         output_parallel = self.quant_method.apply(self, input_parallel, bias_)
+        t2 = _ag2_row_linear_mark(trace_enabled, device) if trace_enabled else 0.0
 
         if self.reduce_results and self.tp_size > 1:
             output = tensor_model_parallel_all_reduce(output_parallel)
         else:
             output = output_parallel
+        t3 = _ag2_row_linear_mark(trace_enabled, device) if trace_enabled else 0.0
+        if trace_enabled:
+            logger.warning(
+                "AG2_ROW_LINEAR_TRACE prefix=%s tokens=%d input_shape=%s "
+                "output_shape=%s tp_size=%d reduce=%s quant=%s split=%.6f "
+                "gemm=%.6f all_reduce=%.6f total=%.6f",
+                self.prefix,
+                int(input_.shape[0]) if input_.dim() > 0 else 0,
+                tuple(input_.shape),
+                tuple(output.shape),
+                self.tp_size,
+                self.reduce_results and self.tp_size > 1,
+                self.quant_method.__class__.__name__,
+                t1 - t0,
+                t2 - t1,
+                t3 - t2,
+                t3 - t0,
+            )
 
         if not self.return_bias:
             return output