[None][feat] Add locality domain Python layer

tensorrt_llm/_torch/cute_dsl_utils.py

@@ -4,11 +4,24 @@
 
 IS_CUTLASS_DSL_AVAILABLE = False
 
+# Whether the public CuTe DSL package provides the SM107/Rubin helper module.
+# Rubin kernels stay disabled when it is absent and callers retain their
+# existing fallback paths.
+IS_CUTLASS_DSL_RUBIN_AVAILABLE = False
+
 if platform.system() != "Windows":
     try:
         import cutlass  # noqa
         import cutlass.cute as cute  # noqa
         logger.info(f"cutlass dsl is available")
         IS_CUTLASS_DSL_AVAILABLE = True
+
+        try:
+            import cutlass.utils.rubin_helpers  # noqa
+        except ImportError:
+            pass
+        else:
+            logger.info("cutlass dsl Rubin helpers are available")
+            IS_CUTLASS_DSL_RUBIN_AVAILABLE = True
     except ImportError:
         pass

tensorrt_llm/_torch/locality_domain/__init__.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+locality domain Localization: Policy and Runtime layers.
+
+This package provides a clean separation of concerns for locality domain execution:
+- layout: Logical/padded shape and partition-slice metadata
+- policy: Planning decisions (enable/disable, partition count, backend selection)
+- runtime: Stream/mempool/event management
+"""
+
+from tensorrt_llm._torch.locality_domain.layout import (
+    PartitionedTensorLayout,
+    make_bf16_linear_output_layout,
+    make_nvfp4_linear_output_layout,
+)
+from tensorrt_llm._torch.locality_domain.policy import (
+    LinearPartitionPlan,
+    LocalityDomainExecutionPlanner,
+    LocalityDomainPolicy,
+    PartitionPlan,
+)
+from tensorrt_llm._torch.locality_domain.runtime import LocalityDomainRuntime
+
+__all__ = [
+    "LocalityDomainPolicy",
+    "PartitionPlan",
+    "LinearPartitionPlan",
+    "PartitionedTensorLayout",
+    "make_bf16_linear_output_layout",
+    "make_nvfp4_linear_output_layout",
+    "LocalityDomainExecutionPlanner",
+    "LocalityDomainRuntime",
+]

tensorrt_llm/_torch/locality_domain/autotune.py

@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Autotuning helpers for kernels that execute on multiple locality domain partitions."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import replace
+from typing import Any
+
+import torch
+
+from tensorrt_llm._torch.autotuner import (
+    AutoTuner,
+    OptimizationProfile,
+    TunableRunner,
+    TuningConfig,
+)
+from tensorrt_llm._torch.locality_domain.runtime import LocalityDomainRuntime
+
+
+class LocalityDomainConcurrentTunableRunner(TunableRunner):
+    """Profile one tactic while every locality domain compute partition is active."""
+
+    def __init__(
+        self,
+        op_runner: TunableRunner,
+        runtime: LocalityDomainRuntime,
+        num_partitions: int,
+        launch_fn: Callable[[int, list[torch.Tensor], object], None],
+    ) -> None:
+        super().__init__()
+        if num_partitions <= 0:
+            raise ValueError(f"num_partitions must be positive, got {num_partitions}")
+        runtime_num_partitions = getattr(runtime, "num_partitions", num_partitions)
+        if runtime_num_partitions != num_partitions:
+            raise ValueError(
+                "num_partitions does not match the locality domain runtime: "
+                f"{num_partitions} != {runtime_num_partitions}"
+            )
+        self._op_runner = op_runner
+        self._runtime = runtime
+        self._num_partitions = num_partitions
+        self._launch_fn = launch_fn
+
+    @property
+    def op_runner(self) -> TunableRunner:
+        """Return the underlying single-partition runner."""
+        return self._op_runner
+
+    def unique_id(self) -> tuple[Any, int, tuple[tuple[int, int], ...]]:
+        return (
+            self._op_runner.unique_id(),
+            self._num_partitions,
+            self._runtime.topology_identity(),
+        )
+
+    def get_valid_tactics(
+        self,
+        inputs: list[torch.Tensor],
+        profile: OptimizationProfile,
+        **kwargs: Any,
+    ) -> list[Any]:
+        return self._op_runner.get_valid_tactics(inputs, profile, **kwargs)
+
+    def should_profile_tactic_in_subprocess(
+        self,
+        custom_op: str,
+        inputs: list[torch.Tensor],
+        tactic: Any,
+        tuning_config: TuningConfig,
+        **kwargs: Any,
+    ) -> bool:
+        # A subprocess cannot reconstruct the process-local CUDA partition
+        # topology, streams, and green contexts used by the launch callback.
+        return False
+
+    def forward(
+        self,
+        inputs: list[torch.Tensor],
+        tactic: Any = -1,
+        **kwargs: Any,
+    ) -> None:
+        self._runtime.fork()
+        try:
+            for partition_id in range(self._num_partitions):
+                with self._runtime.partition_context(partition_id):
+                    self._launch_fn(partition_id, inputs, tactic)
+        finally:
+            self._runtime.join()
+
+
+def tune_locality_domain_concurrent(
+    op_name: str,
+    op_runner: TunableRunner,
+    runtime: LocalityDomainRuntime,
+    num_partitions: int,
+    launch_fn: Callable[[int, list[torch.Tensor], object], None],
+    inputs: list[torch.Tensor],
+    tuning_config: TuningConfig,
+    **choose_one_kwargs: Any,
+) -> tuple[LocalityDomainConcurrentTunableRunner, Any]:
+    """Choose one tactic by profiling all locality domain partitions concurrently."""
+    runner = LocalityDomainConcurrentTunableRunner(
+        op_runner,
+        runtime,
+        num_partitions,
+        launch_fn,
+    )
+    # AutoTuner implements cold-L2 profiling by cloning every tensor into
+    # ordinary CUDA allocations. That loses the VMM node locality of locality domain
+    # weights and would profile a different memory topology. Keep the original
+    # localized tensors, accepting warm-L2 profiling for concurrent locality domain ops.
+    if tuning_config.use_cold_l2_cache:
+        tuning_config = replace(tuning_config, use_cold_l2_cache=False)
+    _, tactic = AutoTuner.get().choose_one(
+        f"{op_name}::locality_domain_concurrent",
+        [runner],
+        tuning_config,
+        inputs,
+        **choose_one_kwargs,
+    )
+    return runner, tactic

tensorrt_llm/_torch/locality_domain/layout.py

@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""locality domain tensor layout metadata.
+
+This module centralizes logical/padded shape and partition-slice reasoning so
+module code does not need to duplicate padding arithmetic.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+NVFP4_WEIGHT_ROW_ALIGNMENT = 32
+BF16_PARTITION_ROW_ALIGNMENT = 8
+
+
+def pad_up(value: int, alignment: int) -> int:
+    """Round value up to the next alignment multiple."""
+    if alignment <= 0:
+        raise ValueError(f"alignment must be positive, got {alignment}")
+    return ((value + alignment - 1) // alignment) * alignment
+
+
+@dataclass(frozen=True)
+class PartitionedTensorLayout:
+    """Logical/padded layout for a tensor split along one axis."""
+
+    logical_shape: tuple[int, ...]
+    padded_shape: tuple[int, ...]
+    partition_axis: int
+    num_partitions: int
+    axis_name: str = "axis"
+    axis_alignment: int = 1
+    alignment_name: str = "alignment"
+
+    def __post_init__(self):
+        rank = len(self.logical_shape)
+        if rank == 0:
+            raise ValueError("logical_shape must not be empty")
+        if rank != len(self.padded_shape):
+            raise ValueError(
+                f"logical_shape rank {rank} does not match padded_shape rank "
+                f"{len(self.padded_shape)}"
+            )
+
+        partition_axis = self.partition_axis
+        if partition_axis < 0:
+            partition_axis += rank
+        if partition_axis < 0 or partition_axis >= rank:
+            raise ValueError(
+                f"partition_axis {self.partition_axis} is out of bounds for rank {rank}"
+            )
+        object.__setattr__(self, "partition_axis", partition_axis)
+
+        if self.num_partitions <= 0:
+            raise ValueError(f"num_partitions must be positive, got {self.num_partitions}")
+        if self.axis_alignment <= 0:
+            raise ValueError(f"axis_alignment must be positive, got {self.axis_alignment}")
+
+        for dim, (logical, padded) in enumerate(zip(self.logical_shape, self.padded_shape)):
+            if logical <= 0:
+                raise ValueError(f"logical_shape[{dim}] must be positive, got {logical}")
+            if padded < logical:
+                raise ValueError(
+                    f"padded_shape[{dim}]={padded} is smaller than logical_shape[{dim}]={logical}"
+                )
+
+    @property
+    def logical_axis_extent(self) -> int:
+        return self.logical_shape[self.partition_axis]
+
+    @property
+    def padded_axis_extent(self) -> int:
+        return self.padded_shape[self.partition_axis]
+
+    @property
+    def is_axis_padding_free(self) -> bool:
+        return self.logical_axis_extent == self.padded_axis_extent
+
+    def per_partition_axis_extent(self, *, padded: bool) -> int:
+        axis_extent = self.padded_axis_extent if padded else self.logical_axis_extent
+        if axis_extent % self.num_partitions != 0:
+            shape_kind = "padded" if padded else "logical"
+            raise ValueError(
+                f"{shape_kind} {self.axis_name}={axis_extent} not divisible by "
+                f"num_partitions={self.num_partitions}"
+            )
+        return axis_extent // self.num_partitions
+
+    def partition_axis_slice(self, partition_id: int, *, padded: bool) -> slice:
+        if partition_id < 0 or partition_id >= self.num_partitions:
+            raise ValueError(
+                f"partition_id={partition_id} out of range for num_partitions={self.num_partitions}"
+            )
+        extent = self.per_partition_axis_extent(padded=padded)
+        return slice(partition_id * extent, (partition_id + 1) * extent)
+
+    def partition_slice(self, partition_id: int, *, padded: bool) -> tuple[slice, ...]:
+        slices = [slice(None)] * len(self.logical_shape)
+        slices[self.partition_axis] = self.partition_axis_slice(partition_id, padded=padded)
+        return tuple(slices)
+
+    def disabled_reason_for_padding_free_split(self) -> str | None:
+        """Return why a split that cannot tolerate per-partition padding is invalid."""
+        if self.logical_axis_extent % self.num_partitions != 0:
+            return (
+                f"{self.axis_name}={self.logical_axis_extent} not divisible "
+                f"by num_partitions={self.num_partitions}"
+            )
+        if self.padded_axis_extent % self.num_partitions != 0:
+            return (
+                f"padded {self.axis_name}={self.padded_axis_extent} not divisible "
+                f"by num_partitions={self.num_partitions}"
+            )
+        if not self.is_axis_padding_free:
+            return (
+                f"{self.axis_name}={self.logical_axis_extent} not divisible "
+                f"by {self.alignment_name}={self.axis_alignment}"
+            )
+        return None
+
+
+def make_nvfp4_linear_output_layout(
+    out_features: int,
+    in_features: int,
+    num_partitions: int,
+    *,
+    padded_out_features: int | None = None,
+) -> PartitionedTensorLayout:
+    """Build layout metadata for NVFP4 Linear partitioning along output N."""
+    if padded_out_features is None:
+        padded_out_features = pad_up(out_features, NVFP4_WEIGHT_ROW_ALIGNMENT)
+    return PartitionedTensorLayout(
+        logical_shape=(out_features, in_features),
+        padded_shape=(padded_out_features, in_features),
+        partition_axis=0,
+        num_partitions=num_partitions,
+        axis_name="out_features",
+        axis_alignment=NVFP4_WEIGHT_ROW_ALIGNMENT,
+        alignment_name="NVFP4 row alignment",
+    )
+
+
+def make_bf16_linear_output_layout(
+    out_features: int,
+    in_features: int,
+    num_partitions: int,
+) -> PartitionedTensorLayout:
+    """Build padding-free BF16 Linear layout metadata for an output-N split.
+
+    The Rubin BF16 kernel requires each partition's contiguous N dimension to
+    be 16-byte aligned. BF16 elements are two bytes, so the full output
+    dimension must be a multiple of ``8 * num_partitions`` elements.
+    """
+    output_row_alignment = BF16_PARTITION_ROW_ALIGNMENT * num_partitions
+    return PartitionedTensorLayout(
+        logical_shape=(out_features, in_features),
+        padded_shape=(
+            pad_up(out_features, output_row_alignment),
+            in_features,
+        ),
+        partition_axis=0,
+        num_partitions=num_partitions,
+        axis_name="out_features",
+        axis_alignment=output_row_alignment,
+        alignment_name="BF16 locality domain row alignment",
+    )