fbgemm-gpu-genai-nightly 2025.12.19__cp310-cp310-manylinux_2_28_x86_64.whl

fbgemm_gpu/runtime_monitor.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import abc
+import logging
+from collections import deque
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Callable, Optional, TypeVar
+
+import torch
+
+
+class TBEStatsReporter(abc.ABC):
+    """
+    Interface for TBE runtime stats reporting. Actual implementation may do
+    custome aggregation (on intended group-key) and reporting destination.
+
+    All the report_XXX functions should be light weighted and fail-safe.
+    """
+
+    @abc.abstractmethod
+    def should_report(self, iteration_step: int) -> bool:
+        """
+        Return whether we should report metrics during this step.
+        This function should be cheap, side-effect free and return immediately.
+        """
+        ...
+
+    @abc.abstractmethod
+    def register_stats(self, stats_name: str, amplifier: int = 1) -> None:
+        """
+        Register stats_name in the whitelist of the reporter
+        """
+        ...
+
+    @abc.abstractmethod
+    def report_duration(
+        self,
+        iteration_step: int,
+        event_name: str,
+        duration_ms: float,
+        embedding_id: str = "",
+        tbe_id: str = "",
+        time_unit: str = "ms",
+        enable_tb_metrics: bool = False,
+    ) -> None:
+        """
+        Report the duration of a timed event.
+        """
+        ...
+
+    @abc.abstractmethod
+    def report_data_amount(
+        self,
+        iteration_step: int,
+        event_name: str,
+        data_bytes: int,
+        embedding_id: str = "",
+        tbe_id: str = "",
+        enable_tb_metrics: bool = False,
+    ) -> None:
+        """
+        Report the size of some data amount.
+        """
+        ...
+
+
+class StdLogStatsReporter(TBEStatsReporter):
+    def __init__(self, report_interval: int) -> None:
+        assert report_interval > 0, "Report interval must be positive"
+        self.report_interval = report_interval
+
+    def register_stats(self, stats_name: str, amplifier: int = 1) -> None:
+        return
+
+    def should_report(self, iteration_step: int) -> bool:
+        return iteration_step % self.report_interval == 0
+
+    def report_duration(
+        self,
+        iteration_step: int,
+        event_name: str,
+        duration_ms: float,
+        embedding_id: str = "",
+        tbe_id: str = "",
+        time_unit: str = "ms",
+        enable_tb_metrics: bool = False,
+    ) -> None:
+        logging.info(
+            f"[Batch #{iteration_step}][TBE:{tbe_id}][Table:{embedding_id}] The event {event_name} took {duration_ms} {time_unit} with {enable_tb_metrics}"
+        )
+
+    def report_data_amount(
+        self,
+        iteration_step: int,
+        event_name: str,
+        data_bytes: int,
+        embedding_id: str = "",
+        tbe_id: str = "",
+        enable_tb_metrics: bool = False,
+    ) -> None:
+        logging.info(
+            f"[Batch #{iteration_step}][TBE:{tbe_id}][Table:{embedding_id}] The event {event_name} used {data_bytes} bytes with {enable_tb_metrics}"
+        )
+
+    def __repr__(self) -> str:
+        return "StdLogStatsReporter{ " f"report_interval={self.report_interval} " "}"
+
+
+@dataclass(frozen=True)
+class TBEStatsReporterConfig:
+    """
+    Configuration for TBEStatsReporter. It eventually instantiates the actual
+    reporter, so it can be deep-copied without incurring the actual reporter
+    getting copied.
+    """
+
+    # Collect required batches every given batches. Non-positive stands for
+    # no collection or reporting
+    interval: int = -1
+
+    def create_reporter(self) -> Optional[TBEStatsReporter]:
+        assert (
+            self.interval <= 0
+        ), "Cannot specify interval without an actual implementation of reporter"
+        return None
+
+
+@dataclass(frozen=True)
+class StdLogStatsReporterConfig(TBEStatsReporterConfig):
+    def create_reporter(self) -> Optional[TBEStatsReporter]:
+        if self.interval <= 0:
+            return None
+        return StdLogStatsReporter(report_interval=self.interval)
+
+
+T = TypeVar("T")
+
+
+class AsyncSeriesTimerRecordedContext:
+    """
+    An easier way to use AsyncSeriesTimer. Example:
+    ```
+    timer : AsyncSeriesTimer = ...
+    with timer.recording(ctx):
+        cuda_kernel1()
+        cuda_kernel2()
+        cuda_kernel3()
+    ```
+    """
+
+    def __init__(
+        self,
+        timer: "AsyncSeriesTimer",
+        context: T,
+        stream: Optional[torch.cuda.Stream] = None,
+    ) -> None:
+        self._context = context
+        self._stream = stream
+        self._timer = timer
+
+    def __enter__(self) -> None:
+        self._timer.start(self._stream)
+
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        self._timer.stop(self._context, self._stream)
+
+
+class AsyncSeriesTimer:
+    """
+    A wrapper class on top of torch.cuda.Event to measure the time between a
+    series of CUDA events. Once initiated, every start() and stop() call pair
+    will measure the timing between them on GPU. Caller cannot initiate another
+    recording if there's already one ongoing.
+
+    Reporting is asynchronous as the timing result is not ready immediately at
+    stop(). Instead, we do it in a lazy way -- we check the all unreported
+    events at every start or stop call.
+    """
+
+    def __init__(self, report_functor: Callable[[T, float], None]) -> None:
+        self._events_queue: deque[tuple[torch.cuda.Event, torch.cuda.Event, T]] = (
+            deque()
+        )
+        self._active_start_event: Optional[torch.cuda.Event] = None
+        self._report_functor = report_functor
+
+    def start(self, stream: Optional[torch.cuda.Stream] = None) -> None:
+        assert self._active_start_event is None, "There's an active recording"
+        self._active_start_event = torch.cuda.Event(enable_timing=True)
+        self._active_start_event.record(stream)
+        self._lazy_report()
+
+    def stop(self, context: T, stream: Optional[torch.cuda.Stream] = None) -> None:
+        assert self._active_start_event is not None, "There's no active recording"
+        active_start_event: torch.cuda.Event = self._active_start_event
+
+        active_stop_event = torch.cuda.Event(enable_timing=True)
+        active_stop_event.record(stream)
+        self._events_queue.append((active_start_event, active_stop_event, context))
+        self._active_start_event = None
+        self._lazy_report()
+
+    def recording(
+        self, context: T, stream: Optional[torch.cuda.Stream] = None
+    ) -> AsyncSeriesTimerRecordedContext:
+        return AsyncSeriesTimerRecordedContext(self, context, stream)
+
+    def _lazy_report(self) -> None:
+        # Since this is a series of timing events, the earliest recorded event
+        # finishes earliest. So we only need to check the leftmost stop event
+        # to decide if we need to report now.
+
+        while len(self._events_queue):
+            stop_event = self._events_queue[0][1]
+            if not stop_event.query():
+                # Even the earliest event hasn't completed in GPU. Don't do
+                # report.
+                return
+            start_event, stop_event, context = self._events_queue.popleft()
+            assert (
+                start_event.query()
+            ), "Recording has start event later than stop event"
+            result = float(start_event.elapsed_time(stop_event))
+            self._report_functor(context, result)

fbgemm_gpu/sll/__init__.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import torch
+
+from fbgemm_gpu.sll.cpu import op_registrations as sll_cpu_registrations
+from fbgemm_gpu.sll.meta import op_registrations as sll_meta_registrations
+from fbgemm_gpu.utils import TorchLibraryFragment
+
+lib = TorchLibraryFragment("fbgemm")
+
+lib.define(
+    """sll_jagged_dense_bmm(
+        Tensor x,
+        Tensor y,
+        Tensor x_offsets,
+        int N,
+        bool allow_tf32,
+        bool use_fbgemm_kernel=True
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_jagged_bmm(
+        Tensor x,
+        Tensor y,
+        Tensor x_offsets,
+        int N,
+        bool allow_tf32,
+        bool use_fbgemm_kernel=True
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_dense_jagged_cat_jagged_out(
+        Tensor a,
+        Tensor b,
+        Tensor a_offsets,
+        int max_seq_len
+    ) -> (Tensor, Tensor)
+    """
+)
+
+lib.define(
+    """sll_jagged_self_substraction_jagged_out(
+        Tensor a,
+        Tensor offsets_a,
+        Tensor offsets_b,
+        int max_seq_len
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged2_to_padded_dense(
+        Tensor values,
+        Tensor offsets,
+        int max_length,
+        float padding_value
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_dense_elementwise_mul_jagged_out(
+        Tensor x,
+        Tensor y,
+        Tensor x_seq_lengths,
+        Tensor x_offsets,
+        int max_seq_len
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_softmax(Tensor x, Tensor x_offsets, int max_seq_len, bool use_fbgemm_kernel=True) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged2_softmax(Tensor x, Tensor offsets, Tensor offsets_total, int max_seq_len, bool transpose) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_array_jagged_bmm_jagged_out(
+        Tensor x,
+        Tensor y,
+        Tensor x_lengths,
+        Tensor x_offsets,
+        Tensor y_lengths,
+        Tensor y_offsets,
+        Tensor z_lengths,
+        Tensor z_offsets,
+        int max_seq_len,
+        bool allow_tf32
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_jagged_bmm_jagged_out(
+        Tensor x,
+        Tensor y,
+        Tensor x_lengths,
+        Tensor x_offsets,
+        Tensor y_lengths,
+        Tensor y_offsets,
+        Tensor z_lengths,
+        Tensor z_offsets,
+        int max_seq_len,
+        bool allow_tf32
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_flash_attention_basic(
+        Tensor q_weights,
+        Tensor k_weights,
+        Tensor v_weights,
+        Tensor offsets,
+        int max_seq_len,
+        bool use_mask=False,
+        bool allow_tf32=True
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_dense_elementwise_add(
+        Tensor x,
+        Tensor x_offsets,
+        Tensor y,
+        int max_seq_len,
+        bool use_fbgemm_kernel=True
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_jagged_dense_flash_attention(
+        Tensor q_weights,
+        Tensor k_weights,
+        Tensor v_weights,
+        Tensor attn_bias,
+        Tensor offsets,
+        int max_seq_len,
+        bool allow_tf32=True
+    ) -> Tensor
+    """
+)
+
+lib.define(
+    """sll_multi_head_jagged_flash_attention(
+        Tensor q_weights,
+        Tensor k_weights,
+        Tensor v_weights,
+        Tensor offsets,
+        int max_seq_len,
+        bool allow_tf32=True
+    ) -> Tensor
+    """
+)
+
+# NOTE: here we register the op for AutogradCUDA/CPU and CUDA/CPU with the same
+# function however, this is not ideal because in the inference case, we don't
+# need the autograd forward to save the context because we don't need to do
+# backward.
+
+for op_name, dispatches in sll_cpu_registrations.items():
+    lib.register(op_name, dispatches)
+
+for op_name, dispatches in sll_meta_registrations.items():
+    lib.register(op_name, dispatches)
+
+if torch.cuda.is_available():
+    from fbgemm_gpu.sll.triton import op_registrations as sll_gpu_registrations
+
+    for op_name, dispatches in sll_gpu_registrations.items():
+        lib.register(op_name, dispatches)

fbgemm_gpu/sll/cpu/__init__.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from fbgemm_gpu.sll.cpu.cpu_sll import (  # noqa F401
+    cpu_array_jagged_bmm_jagged_out,
+    cpu_array_jagged_bmm_jagged_out_kernel,  # noqa F401
+    cpu_dense_jagged_cat_jagged_out,
+    cpu_jagged2_softmax,
+    cpu_jagged2_to_padded_dense,
+    cpu_jagged_dense_bmm,
+    cpu_jagged_dense_elementwise_add,
+    cpu_jagged_dense_elementwise_mul_jagged_out,
+    cpu_jagged_dense_flash_attention,
+    cpu_jagged_flash_attention_basic,
+    cpu_jagged_jagged_bmm,
+    cpu_jagged_jagged_bmm_jagged_out,
+    cpu_jagged_jagged_bmm_jagged_out_kernel,  # noqa F401
+    cpu_jagged_self_substraction_jagged_out,
+    cpu_jagged_softmax,
+)
+
+# pyre-ignore[5]
+op_registrations = {
+    "sll_jagged_dense_bmm": {
+        "CPU": cpu_jagged_dense_bmm,
+        "AutogradCPU": cpu_jagged_dense_bmm,
+    },
+    "sll_jagged_jagged_bmm": {
+        "CPU": cpu_jagged_jagged_bmm,
+        "AutogradCPU": cpu_jagged_jagged_bmm,
+    },
+    "sll_dense_jagged_cat_jagged_out": {
+        "CPU": cpu_dense_jagged_cat_jagged_out,
+    },
+    "sll_jagged_self_substraction_jagged_out": {
+        "CPU": cpu_jagged_self_substraction_jagged_out,
+    },
+    "sll_jagged2_to_padded_dense": {
+        "CPU": cpu_jagged2_to_padded_dense,
+        "AutogradCPU": cpu_jagged2_to_padded_dense,
+    },
+    "sll_jagged_dense_elementwise_mul_jagged_out": {
+        "CPU": cpu_jagged_dense_elementwise_mul_jagged_out,
+        "AutogradCPU": cpu_jagged_dense_elementwise_mul_jagged_out,
+    },
+    "sll_jagged_softmax": {
+        "CPU": cpu_jagged_softmax,
+        "AutogradCPU": cpu_jagged_softmax,
+    },
+    "sll_jagged2_softmax": {
+        "CPU": cpu_jagged2_softmax,
+        "AutogradCPU": cpu_jagged2_softmax,
+    },
+    "sll_array_jagged_bmm_jagged_out": {
+        "CPU": cpu_array_jagged_bmm_jagged_out,
+        "AutogradCPU": cpu_array_jagged_bmm_jagged_out,
+    },
+    "sll_jagged_jagged_bmm_jagged_out": {
+        "CPU": cpu_jagged_jagged_bmm_jagged_out,
+        "AutogradCPU": cpu_jagged_jagged_bmm_jagged_out,
+    },
+    "sll_jagged_flash_attention_basic": {
+        "CPU": cpu_jagged_flash_attention_basic,
+        "AutogradCPU": cpu_jagged_flash_attention_basic,
+    },
+    "sll_jagged_dense_elementwise_add": {
+        "CPU": cpu_jagged_dense_elementwise_add,
+        "AutogradCPU": cpu_jagged_dense_elementwise_add,
+    },
+    "sll_jagged_dense_flash_attention": {
+        "CPU": cpu_jagged_dense_flash_attention,
+        "AutogradCPU": cpu_jagged_dense_flash_attention,
+    },
+}