[None][feat] Wire Rubin locality-domain execution through PyTorch models

tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py

@@ -1096,6 +1096,7 @@ def forward_sparse_attn(
         and os.environ.get("TRTLLM_MLA_Q_B_PROJ_USE_CUTE_DSL", "1") == "1"
         and self.q_b_proj.bias is None
         and self.q_b_proj.weight.dtype == torch.bfloat16
+        and getattr(self.q_b_proj, "_locality_domain_weight_shards", None) is None
         and not _is_fused_q_fp8_quant_enabled(
             self, num_generations=num_generations, num_contexts=num_contexts
         )

tensorrt_llm/_torch/attention/mla.py

@@ -30,6 +30,8 @@
 
 from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE
 from ..distributed import AllReduceParams
+from ..locality_domain.policy import LocalityDomainExecutionPlanner, LocalityDomainPolicy
+from ..locality_domain.runtime import LocalityDomainRuntime
 from ..model_config import ModelConfig
 from ..modules.linear import Linear, TensorParallelMode, is_static_nvfp4_input_eligible
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
@@ -251,6 +253,7 @@ def __init__(
         o_lora_rank: int = 1024,
         fuse_qkv_a_proj: bool = True,
         rms_norm_eps: Optional[float] = None,
+        enable_locality_domain_bf16_linear: bool = False,
     ) -> None:
         """
         Initialize the MLA module.
@@ -441,6 +444,15 @@ def __init__(
             None
         )
         self.use_cute_dsl_bf16_gemm = config.use_cute_dsl_bf16_gemm
+        locality_domain_policy = config.extra_attrs.get("locality_domain_policy")
+        self._locality_domain_policy = (
+            locality_domain_policy if locality_domain_policy is not None else LocalityDomainPolicy()
+        )
+        self._locality_domain_runtime: Optional[LocalityDomainRuntime] = None
+        self._locality_domain_k_b_proj_trans_shards: Optional[tuple[torch.Tensor, torch.Tensor]] = (
+            None
+        )
+        self._locality_domain_v_b_proj_shards: Optional[tuple[torch.Tensor, torch.Tensor]] = None
 
         if not self.is_lite:
             kv_a_out_features = self.kv_lora_rank + self.qk_rope_head_dim
@@ -490,6 +502,8 @@ def __init__(
                 force_dynamic_quantization=config.force_dynamic_quantization,
                 use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
                 use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm,
+                enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear,
+                locality_domain_policy=self._locality_domain_policy,
             )
         else:
             self.kv_a_proj_with_mqa = Linear(
@@ -518,6 +532,8 @@ def __init__(
                 force_dynamic_quantization=config.force_dynamic_quantization,
                 use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
                 use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm,
+                enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear,
+                locality_domain_policy=self._locality_domain_policy,
             )
             self.q_b_proj = self.q_proj
 
@@ -568,6 +584,8 @@ def yarn_get_mscale(scale=1, mscale=1):
                 force_dynamic_quantization=config.force_dynamic_quantization,
                 use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
                 use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm,
+                enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear,
+                locality_domain_policy=self._locality_domain_policy,
             )
             self.v_b_proj = nn.Parameter(
                 torch.empty(
@@ -652,6 +670,8 @@ def yarn_get_mscale(scale=1, mscale=1):
                 force_dynamic_quantization=config.force_dynamic_quantization,
                 use_cute_dsl_blockscaling_mm=self.use_cute_dsl_blockscaling_mm,
                 use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm,
+                enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear,
+                locality_domain_policy=self._locality_domain_policy,
             )
             self.attention_output_hidden_size = self.o_proj.in_features
 
@@ -776,6 +796,8 @@ def create_weights(self):
         else:
             self.k_b_proj_trans_scale = None
             self.v_b_proj_scale = None
+        self._locality_domain_k_b_proj_trans_shards = None
+        self._locality_domain_v_b_proj_shards = None
         self._weights_transformed = False
         self._init_fused_q_fp8_context()
 
@@ -1533,9 +1555,19 @@ def forward_context(
             q, compressed_kv, k_pe, position_ids, attn_metadata, output, latent_cache
         )
 
-    def _bmm_bf16_out(self, a, b_no_transpose, b_transposed, output):
+    def _bmm_bf16_out(
+        self, a, b_no_transpose, b_transposed, output, locality_domain_weight_shards=None
+    ):
         """BMM with optional CuTe DSL bf16 acceleration on Blackwell/SM107."""
-        if self.use_cute_dsl_bf16_bmm and is_sm_100f():
+        if (
+            self.use_cute_dsl_bf16_bmm
+            and get_sm_version() == 107
+            and locality_domain_weight_shards is not None
+        ):
+            torch.ops.trtllm.cute_dsl_bf16_bmm_locality_domain_inplace_rubin(
+                a, locality_domain_weight_shards[0], locality_domain_weight_shards[1], output
+            )
+        elif self.use_cute_dsl_bf16_bmm and is_sm_100f():
             bf16_bmm_op = (
                 torch.ops.trtllm.cute_dsl_bf16_bmm_rubin
                 if get_sm_version() == 107
@@ -1683,6 +1715,7 @@ def _mla_gen_rope():
                         self.k_b_proj_trans,
                         self.k_b_proj_trans.transpose(1, 2),
                         q_nope_out,
+                        self._locality_domain_k_b_proj_trans_shards,
                     ),
                     _mla_gen_rope,
                     self.ln_events[0],
@@ -1762,6 +1795,7 @@ def _mla_gen_rope():
                 self.v_b_proj,
                 self.v_b_proj.transpose(1, 2),
                 attn_output.transpose(0, 1),
+                self._locality_domain_v_b_proj_shards,
             )
         elif self.v_b_proj.dtype == torch.float8_e4m3fn:
             fp8_block_scaling_bmm_out(
@@ -1901,6 +1935,7 @@ def forward_absorption_context(
                     self.k_b_proj_trans,
                     self.k_b_proj_trans.transpose(1, 2),
                     q_nope_out,
+                    self._locality_domain_k_b_proj_trans_shards,
                 )
             elif self.k_b_proj_trans.dtype == torch.float8_e4m3fn:
                 # [num_heads, num_tokens, self.kv_lora_rank]
@@ -1966,6 +2001,7 @@ def forward_absorption_context(
                 self.v_b_proj,
                 self.v_b_proj.transpose(1, 2),
                 attn_output.transpose(0, 1),
+                self._locality_domain_v_b_proj_shards,
             )
         elif self.v_b_proj.dtype == torch.float8_e4m3fn:
             fp8_block_scaling_bmm_out(
@@ -2167,10 +2203,107 @@ def transform_weights(self) -> None:
             self.v_b_proj, self.v_b_proj_scale = self.resmooth_parameters(
                 self.v_b_proj, self.v_b_proj_scale, recipe=(1, 128, 128)
             )
+        self._rebuild_locality_domain_bf16_bmm_weight_shards()
         self._weights_transformed = True
 
     def cache_derived_state(self) -> None:
         self._weights_transformed = True
 
     def post_load_weights(self) -> None:
         self.transform_weights()
+        self.cache_derived_state()
+
+    def _bind_v_b_proj_weight(self, weight: torch.Tensor) -> None:
+        """Bind the V absorption weight without retaining a released full weight."""
+        plan = self.kv_b_proj.partition_plan
+        if plan.enabled and plan.op_kind == "bf16_linear":
+            weight = weight.clone(memory_format=torch.contiguous_format)
+        self.v_b_proj = nn.Parameter(weight, requires_grad=False)
+        self._locality_domain_k_b_proj_trans_shards = None
+        self._locality_domain_v_b_proj_shards = None
+        self._weights_transformed = False
+
+    def _has_full_absorption_weights(self) -> bool:
+        if getattr(self, "kv_lora_rank", None) != 512 or getattr(self, "v_head_dim", None) != 128:
+            return False
+        # The Rubin BF16 BMM requires its K dimension to be 16-byte aligned.
+        # qk_nope_head_dim is K for the absorbed K projection.
+        if getattr(self, "qk_nope_head_dim", 0) % 8 != 0:
+            return False
+
+        k_weight = getattr(self, "k_b_proj_trans", None)
+        v_weight = getattr(self, "v_b_proj", None)
+        if not isinstance(k_weight, torch.Tensor) or not isinstance(v_weight, torch.Tensor):
+            return False
+        if k_weight.is_meta or v_weight.is_meta:
+            return False
+
+        expected_k_shape = (
+            self.num_heads_tp,
+            self.kv_lora_rank,
+            self.qk_nope_head_dim,
+        )
+        expected_v_shape = (
+            self.num_heads_tp_cp,
+            self.v_head_dim,
+            self.kv_lora_rank,
+        )
+        return (
+            k_weight.dtype == v_weight.dtype
+            and tuple(k_weight.shape) == expected_k_shape
+            and tuple(v_weight.shape) == expected_v_shape
+            and k_weight.numel() == math.prod(expected_k_shape)
+            and v_weight.numel() == math.prod(expected_v_shape)
+        )
+
+    def _copy_absorption_weight_shards(
+        self, weight: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        assert self._locality_domain_runtime is not None
+        num_partitions = self._locality_domain_runtime.num_partitions
+        assert weight.shape[1] % num_partitions == 0
+
+        shard_n = weight.shape[1] // num_partitions
+        shards = []
+        for partition_id in range(num_partitions):
+            weight_slice = weight[
+                :,
+                partition_id * shard_n : (partition_id + 1) * shard_n,
+                :,
+            ]
+            with self._locality_domain_runtime.partition_weight_context(partition_id):
+                shard = torch.empty_like(
+                    weight_slice,
+                    memory_format=torch.contiguous_format,
+                )
+                shard.copy_(weight_slice)
+            shards.append(shard)
+        assert len(shards) == 2
+        return shards[0], shards[1]
+
+    def _rebuild_locality_domain_bf16_bmm_weight_shards(self) -> None:
+        """Rebuild localized BF16 absorption weights from the full parameters."""
+        self._locality_domain_k_b_proj_trans_shards = None
+        self._locality_domain_v_b_proj_shards = None
+
+        if not self._has_full_absorption_weights():
+            return
+
+        locality_domain_policy = getattr(self, "_locality_domain_policy", None)
+        if locality_domain_policy is None:
+            return
+        plan = LocalityDomainExecutionPlanner(locality_domain_policy).plan_bf16_bmm(
+            dtype=self.k_b_proj_trans.dtype,
+            use_cute_dsl_bf16_bmm=getattr(self, "use_cute_dsl_bf16_bmm", False),
+        )
+        if not plan.enabled:
+            return
+
+        if getattr(self, "_locality_domain_runtime", None) is None:
+            self._locality_domain_runtime = LocalityDomainRuntime(plan.num_partitions)
+
+        k_shards = self._copy_absorption_weight_shards(self.k_b_proj_trans)
+        v_shards = self._copy_absorption_weight_shards(self.v_b_proj)
+        self._locality_domain_runtime.prepare_for_capture(plan)
+        self._locality_domain_k_b_proj_trans_shards = k_shards
+        self._locality_domain_v_b_proj_shards = v_shards

tensorrt_llm/_torch/compilation/utils.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import contextlib
 from typing import Callable, List, Optional, Union
 
@@ -248,6 +251,12 @@ def inplace_info():
         "cute_dsl_bf16_gemm_rubin": {
             1: "output"
         },
+        "cute_dsl_bf16_gemm_locality_domain_inplace_rubin": {
+            1: "output"
+        },
+        "cute_dsl_bf16_bmm_locality_domain_inplace_rubin": {
+            1: "output"
+        },
         "cute_dsl_fp8_bmm_rubin": {
             1: "output"
         },