diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py index b50945c23056..0da25e436d7d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py +++ b/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 ) diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index eb6ad4f97bc9..b20229b73cbb 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/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,6 +2203,7 @@ 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: @@ -2174,3 +2211,99 @@ def cache_derived_state(self) -> None: 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 diff --git a/tensorrt_llm/_torch/compilation/utils.py b/tensorrt_llm/_torch/compilation/utils.py index 2e98607b84b6..2c7866c34707 100644 --- a/tensorrt_llm/_torch/compilation/utils.py +++ b/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" }, diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index e641111adb4e..5237ac56c277 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -10381,6 +10381,12 @@ def forward( a_tensor, b_tensor, c_tensor = inputs batch_size, m, k = a_tensor.shape n = b_tensor.shape[1] + locality_domain_id = get_current_locality_domain() + if locality_domain_id is not None: + assert locality_domain_id in (0, 1) + assert tuple(c_tensor.shape) == (batch_size, m, 2 * n) + c_tensor = c_tensor[:, :, locality_domain_id * + n:(locality_domain_id + 1) * n] # C is passed as [M, N, B]; from_dlpack captures the real strides so # non-contiguous views are written by TMA without a copy. @@ -10508,6 +10514,264 @@ def forward( stream=stream, ) + def _validate_cute_dsl_bf16_bmm_locality_domain_inputs( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + ) -> None: + tensors = (input, weight_0, weight_1, output) + if any(tensor.dim() != 3 for tensor in tensors): + raise ValueError( + "locality domain BF16 BMM input, weight shards, and output must be 3-D, " + f"got shapes {[tuple(tensor.shape) for tensor in tensors]}.") + if input.dtype != torch.bfloat16: + raise ValueError( + "locality domain BF16 BMM input must have dtype bfloat16, " + f"got {input.dtype}.") + if weight_0.dtype != torch.bfloat16 or weight_1.dtype != torch.bfloat16: + raise ValueError( + "locality domain BF16 BMM weight shards must have dtype bfloat16, " + f"got {weight_0.dtype} and {weight_1.dtype}.") + if output.dtype != torch.bfloat16: + raise ValueError( + "locality domain BF16 BMM output must have dtype bfloat16, " + f"got {output.dtype}.") + if weight_0.shape != weight_1.shape: + raise ValueError( + "locality domain BF16 BMM weight shards must have identical shapes, " + f"got {tuple(weight_0.shape)} and {tuple(weight_1.shape)}.") + if input.shape[0] != weight_0.shape[0]: + raise ValueError( + "locality domain BF16 BMM input and weight shards must have the same " + f"batch size, got {input.shape[0]} and {weight_0.shape[0]}.") + if input.shape[2] != weight_0.shape[2]: + raise ValueError( + "locality domain BF16 BMM input and weight shards must have the same K " + f"dimension, got {input.shape[2]} and {weight_0.shape[2]}.") + alignment_elements = 16 // input.element_size() + if (input.shape[2] % alignment_elements != 0 + or weight_0.shape[1] % alignment_elements != 0): + raise ValueError( + "locality domain BF16 BMM K and per-shard N dimensions must be " + f"multiples of {alignment_elements} for 16-byte alignment, " + f"got K={input.shape[2]} and N={weight_0.shape[1]}.") + expected_output_shape = ( + input.shape[0], + input.shape[1], + weight_0.shape[1] + weight_1.shape[1], + ) + if output.shape != expected_output_shape: + raise ValueError( + "locality domain BF16 BMM output shape must be [B, M, N0 + N1], " + f"expected {expected_output_shape}, got {tuple(output.shape)}.") + if any(tensor.device != input.device for tensor in tensors[1:]): + raise ValueError( + "locality domain BF16 BMM input, weight shards, and output must be on " + f"the same device, got {[tensor.device for tensor in tensors]}." + ) + for tensor_name, tensor in zip( + ("input", "weight_0", "weight_1", "output"), tensors): + _validate_16_byte_aligned_dense_tensor(tensor, tensor_name) + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_bmm_locality_domain_inplace_rubin", + mutates_args=("output", ), + device_types="cuda", + schema="(Tensor input, Tensor weight_0, Tensor weight_1, " + "Tensor(a!) output, bool use_tvm_ffi=True) -> ()", + ) + @_with_input_cuda_device + def cute_dsl_bf16_bmm_locality_domain_inplace_rubin( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + """Tune and launch both Rubin locality domain BF16 BMM partitions. + + Call once before the first CUDA graph capture to initialize locality domain + resources and compile the selected partition-local artifacts. + """ + if get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + raise ValueError( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + "CuteDSL BF16 BMM Rubin only supports SM107 with " + "nvidia-cutlass-dsl with SM107 support.") + _validate_cute_dsl_bf16_bmm_locality_domain_inputs( + input, weight_0, weight_1, output) + + runtime = LocalityDomainRuntime(num_partitions=2) + op_runner = CuteDSLBf16RubinBmmRunner(use_tvm_ffi=use_tvm_ffi) + inputs = [input, weight_0, output] + + def launch_partition( + partition_id: int, + partition_inputs: List[torch.Tensor], + tactic, + ) -> None: + weight = partition_inputs[1] if partition_id == 0 else weight_1 + op_runner( + [partition_inputs[0], weight, partition_inputs[2]], + tactic=tactic, + ) + + runner, best_tactic = tune_locality_domain_concurrent( + "trtllm::cute_dsl_bf16_bmm_locality_domain_inplace_rubin", + op_runner, + runtime, + 2, + launch_partition, + inputs, + op_runner.__class__.tuning_config, + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_bmm_locality_domain_inplace_rubin") + def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + return None + + def _validate_cute_dsl_bf16_gemm_locality_domain_inputs( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + ) -> None: + tensors = (input, weight_0, weight_1, output) + if any(tensor.dim() != 2 for tensor in tensors): + raise ValueError( + "locality domain BF16 GEMM input, weight shards, and output must be 2-D, " + f"got shapes {[tuple(tensor.shape) for tensor in tensors]}.") + if input.dtype != torch.bfloat16: + raise ValueError( + "locality domain BF16 GEMM input must have dtype bfloat16, " + f"got {input.dtype}.") + if weight_0.dtype != torch.bfloat16 or weight_1.dtype != torch.bfloat16: + raise ValueError( + "locality domain BF16 GEMM weight shards must have dtype bfloat16, " + f"got {weight_0.dtype} and {weight_1.dtype}.") + if output.dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + "locality domain BF16 GEMM output must have dtype bfloat16 or float32, " + f"got {output.dtype}.") + if weight_0.shape != weight_1.shape: + raise ValueError( + "locality domain BF16 GEMM weight shards must have identical shapes, " + f"got {tuple(weight_0.shape)} and {tuple(weight_1.shape)}.") + if input.shape[1] != weight_0.shape[1]: + raise ValueError( + "locality domain BF16 GEMM input and weight shards must have the same K " + f"dimension, got {input.shape[1]} and {weight_0.shape[1]}.") + k_alignment_elements = 16 // input.element_size() + n_alignment_elements = 16 // output.element_size() + if (input.shape[1] % k_alignment_elements != 0 + or weight_0.shape[0] % n_alignment_elements != 0): + raise ValueError( + "locality domain BF16 GEMM K and per-shard N dimensions must be " + "multiples of their respective 16-byte alignments " + f"(K: {k_alignment_elements}, N: {n_alignment_elements}), " + f"got K={input.shape[1]} and N={weight_0.shape[0]}.") + expected_output_shape = ( + input.shape[0], + weight_0.shape[0] + weight_1.shape[0], + ) + if output.shape != expected_output_shape: + raise ValueError( + "locality domain BF16 GEMM output shape must be [M, N0 + N1], " + f"expected {expected_output_shape}, got {tuple(output.shape)}.") + if any(tensor.device != input.device for tensor in tensors[1:]): + raise ValueError( + "locality domain BF16 GEMM input, weight shards, and output must be on " + f"the same device, got {[tensor.device for tensor in tensors]}." + ) + for tensor_name, tensor in ( + ("input", input), + ("weight_0", weight_0), + ("weight_1", weight_1), + ): + # The GEMM runner canonicalizes non-contiguous A/B tensors. A + # contiguous view is reused directly and must honor its declared + # 16-byte pointer alignment. + if tensor.is_contiguous() and tensor.data_ptr() % 16 != 0: + raise ValueError( + f"{tensor_name} data pointer must be 16-byte aligned, got " + f"data_ptr={tensor.data_ptr()}.") + _validate_16_byte_aligned_dense_tensor(output, "output") + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_gemm_locality_domain_inplace_rubin", + mutates_args=("output", ), + device_types="cuda", + schema="(Tensor input, Tensor weight_0, Tensor weight_1, " + "Tensor(a!) output, bool use_tvm_ffi=True) -> ()", + ) + @_with_input_cuda_device + def cute_dsl_bf16_gemm_locality_domain_inplace_rubin( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + """Tune and launch both Rubin locality domain BF16 GEMM partitions. + + Call once before the first CUDA graph capture to initialize locality domain + resources and compile the selected partition-local artifacts. + """ + if get_sm_version() != 107 or not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + raise ValueError( + f"CuteDSL: SM version {get_sm_version()} is not supported. " + "CuteDSL BF16 GEMM Rubin only supports SM107 with " + "nvidia-cutlass-dsl with SM107 support.") + _validate_cute_dsl_bf16_gemm_locality_domain_inputs( + input, weight_0, weight_1, output) + + runtime = LocalityDomainRuntime(num_partitions=2) + op_runner = CuteDSLBf16RubinGemmRunner(use_tvm_ffi=use_tvm_ffi, + output_dtype=output.dtype) + inputs = [input, weight_0, output] + + def launch_partition( + partition_id: int, + partition_inputs: List[torch.Tensor], + tactic, + ) -> None: + weight = partition_inputs[1] if partition_id == 0 else weight_1 + op_runner( + [partition_inputs[0], weight, partition_inputs[2]], + tactic=tactic, + ) + + runner, best_tactic = tune_locality_domain_concurrent( + "trtllm::cute_dsl_bf16_gemm_locality_domain_inplace_rubin", + op_runner, + runtime, + 2, + launch_partition, + inputs, + op_runner.__class__.tuning_config, + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_gemm_locality_domain_inplace_rubin") + def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> None: + return None + @torch.library.custom_op("trtllm::cute_dsl_bf16_bmm_rubin", mutates_args=("output", ), device_types="cuda") @@ -10631,14 +10895,21 @@ def forward( m, k = a_tensor.shape n = b_tensor.shape[0] batch_size = 1 - if split_k_slices > 1 and tuple(c_tensor.shape) != (m, n): + locality_domain_id = get_current_locality_domain() + if locality_domain_id is not None: + assert locality_domain_id in (0, 1) + assert tuple(c_tensor.shape) == (m, 2 * n) + c_tensor = c_tensor[:, locality_domain_id * + n:(locality_domain_id + 1) * n] + elif split_k_slices > 1 and tuple(c_tensor.shape) != (m, n): raise RuntimeError( "BF16 split-K GEMM requires an [M, N] output, got " f"output.shape={tuple(c_tensor.shape)}, expected={(m, n)}.") a_tensor = a_tensor.contiguous() b_tensor = b_tensor.contiguous() - c_needs_copy = not c_tensor.is_contiguous() + c_needs_copy = locality_domain_id is None and not c_tensor.is_contiguous( + ) c_buf = torch.empty_like(c_tensor) if c_needs_copy else c_tensor a_batched = a_tensor.unsqueeze(0) # [1, M, K] @@ -14110,6 +14381,106 @@ def _( return input.new_empty((input.shape[0], weight.shape[0]), dtype=torch.float8_e4m3fn) + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_gemm_rubin", + mutates_args=(), + device_types="cuda", + schema= + "(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, " + "ScalarType output_dtype, bool to_userbuffers=False, bool use_tvm_ffi=True, " + "Tensor? output_tensor=None, SymInt partition_id=-1) -> Tensor?") + def cute_dsl_nvfp4_gemm_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + to_userbuffers: bool = False, + use_tvm_ffi: bool = True, + output_tensor: Optional[torch.Tensor] = None, + partition_id: int = -1, + ) -> Optional[torch.Tensor]: + """CuteDSL-based NVFP4 GEMM optimized for Rubin (SM107). + + Args: + input: Activation tensor [m, k] in FP4 format (packed in uint8) + weight: Weight tensor [n, k] in FP4 format (packed in uint8) + input_scale: Activation scale factors + weight_scale: Weight scale factors + alpha: Scaling factor + output_tensor: Deprecated. Use cute_dsl_nvfp4_gemm_inplace_rubin for locality domain inplace mode. + output_dtype: Output data type (must be bfloat16) + to_userbuffers: Whether to allocate output from UserBuffers pool + use_tvm_ffi: Whether to use TVM-FFI to call the kernel. Enable this option could help reduce the kernel host launch overhead. + partition_id: Must be -1 for this non-mutating op. + + Note: + This function is primarily used internally by nvfp4_gemm. + Direct usage is discouraged. Consider using nvfp4_gemm instead + for automatic backend selection with better performance. + """ + # Validate SM version before attempting to use CuteDSL + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"CuteDSL NVFP4 Rubin backend requires SM 107, but got SM {sm_version}. " + f"Please use nvfp4_gemm with backend='auto' for automatic backend selection." + ) + + tuner = AutoTuner.get() + + if output_tensor is None: + if partition_id != -1: + raise ValueError( + "partition_id must be -1 when output_tensor is not provided." + ) + runner = CuteDSLNVFP4RubinLinear(output_dtype, to_userbuffers, + use_tvm_ffi) + inputs = [input, weight, input_scale, weight_scale, alpha] + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_nvfp4_gemm_rubin", + [runner], + runner.__class__.tuning_config, + inputs, + ) + return runner(inputs, tactic=best_tactic) + else: + raise ValueError( + "output_tensor mode mutates its output and must use " + "torch.ops.trtllm.cute_dsl_nvfp4_gemm_inplace_rubin instead." + ) + + @torch.library.register_fake("trtllm::cute_dsl_nvfp4_gemm_rubin") + def _( + mat_a: torch.Tensor, + mat_b: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, # Match custom op signature + output_dtype: torch.dtype, + to_userbuffers: bool = False, + use_tvm_ffi: bool = True, + output_tensor: Optional[torch.Tensor] = None, + partition_id: int = -1, + ) -> Optional[torch.Tensor]: + # [m, k] + shape = list(mat_a.shape) + # [n, k] + shape[-1] = mat_b.shape[-2] + # output is fixed as bf16 + ret = mat_a.new_empty(shape, dtype=torch.bfloat16) + if output_tensor is None: + if partition_id != -1: + raise ValueError( + "partition_id must be -1 when output_tensor is not provided." + ) + return ret + else: + raise ValueError( + "output_tensor mode mutates its output and must use " + "torch.ops.trtllm.cute_dsl_nvfp4_gemm_inplace_rubin instead." + ) + @torch.library.custom_op( "trtllm::cute_dsl_nvfp4_gemm_inplace_rubin", mutates_args=("output_tensor", ), diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv3.py b/tensorrt_llm/_torch/models/modeling_deepseekv3.py index 682e530eed1c..6a54201f9065 100755 --- a/tensorrt_llm/_torch/models/modeling_deepseekv3.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv3.py @@ -58,6 +58,7 @@ from ..attention.mla import MLA from ..distributed import (AllReduce, AllReduceFusionOp, AllReduceParams, MoEAllReduce, MoEAllReduceParams, allgather) +from ..locality_domain.policy import LocalityDomainPolicy from ..model_config import ModelConfig from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding @@ -407,8 +408,7 @@ def detect_shared_mtp_weights() -> bool: attn_module = all_named_modules[parent_module_name] _, v_b_proj = split_kv_b_proj(module.weight.data, is_scale=False) - attn_module.v_b_proj = nn.Parameter(v_b_proj, - requires_grad=False) + attn_module._bind_v_b_proj_weight(v_b_proj) attn_module.k_b_proj_trans.data.copy_( k_b_proj_trans.reshape( @@ -717,6 +717,7 @@ def __init__( use_cute_dsl_bf16_gemm: bool = False, use_cute_dsl_blockscaling_mm: bool = False, lora: Optional[LoraLayer] = None, + locality_domain_policy: Optional[LocalityDomainPolicy] = None, ): super().__init__( in_features, @@ -732,6 +733,8 @@ def __init__( skip_create_weights_in_init=skip_create_weights_in_init, use_custom_cublas_mm=use_custom_cublas_mm, use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + enable_locality_domain_bf16_linear=True, + locality_domain_policy=locality_domain_policy, lora=lora, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, ) @@ -741,12 +744,12 @@ def apply_linear(self, bias, lora_params: Optional[dict] | None = None, layer_idx: Optional[int] | None = None): + if self._locality_domain_weight_shards is not None: + return super().apply_linear(input, bias, lora_params, layer_idx) num_tokens = input.shape[0] has_any_quant = self.has_any_quant - use_cute_dsl_bf16_gemm = (self.use_cute_dsl_bf16_gemm - and not has_any_quant and is_sm_100f() - and self.weight.dtype == torch.bfloat16) - if (not use_cute_dsl_bf16_gemm and not has_any_quant + # The fused op has no LoRA inputs; let Linear merge active adapters. + if (not has_any_quant and (self.lora is None or not lora_params) and 1 <= num_tokens <= 16 and get_sm_version() not in [120, 121]): output = torch.ops.trtllm.dsv3_fused_a_gemm_op( @@ -790,7 +793,8 @@ def __init__( config=model_config, aux_stream_dict=aux_stream_dict, mapping_with_cp=mapping_with_cp, - reduce_output=reduce_output) + reduce_output=reduce_output, + enable_locality_domain_bf16_linear=True) self.kv_a_proj_with_mqa = DeepseekV3Linear( config.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim + @@ -804,6 +808,7 @@ def __init__( use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm, use_cute_dsl_blockscaling_mm=model_config. use_cute_dsl_blockscaling_mm, + locality_domain_policy=model_config.locality_domain_policy, ) @@ -842,7 +847,8 @@ def __init__( config=model_config, aux_stream_dict=aux_stream_dict, mapping_with_cp=mapping_with_cp, - reduce_output=reduce_output) + reduce_output=reduce_output, + enable_locality_domain_bf16_linear=True) self.kv_a_proj_with_mqa = DeepseekV3Linear( config.hidden_size, @@ -853,10 +859,11 @@ def __init__( skip_create_weights_in_init=model_config. skip_create_weights_in_init, use_custom_cublas_mm=True, - use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm) + use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm, + locality_domain_policy=model_config.locality_domain_policy) -class DeepseekV3Gate(nn.Module): +class DeepseekV3Gate(Linear): def __init__( self, @@ -871,12 +878,19 @@ def __init__( apply_routing: bool = False, moe_backend: str = 'CUTLASS', use_cute_dsl_bf16_gemm: bool = False, + locality_domain_policy: Optional[LocalityDomainPolicy] = None, ): - super().__init__() - self.use_cute_dsl_bf16_gemm = use_cute_dsl_bf16_gemm - self.weight = nn.Parameter(torch.empty((num_experts, hidden_size), - dtype=dtype), - requires_grad=False) + super().__init__( + hidden_size, + num_experts, + bias=False, + dtype=dtype, + reduce_output=False, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + enable_locality_domain_bf16_linear=True, + locality_domain_policy=locality_domain_policy, + ) + self._skip_layerwise_quant_config = True self.moe_backend = moe_backend if moe_backend == 'TRTLLM': bias_dtype = torch.bfloat16 @@ -897,8 +911,11 @@ def __init__( is_fused=fuse_routing_kernel) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if (self.use_cute_dsl_bf16_gemm and is_sm_100f() - and self.weight.dtype == torch.bfloat16): + if self._locality_domain_weight_shards is not None: + logits = self._run_bf16_linear_locality_domain( + hidden_states, None, output_dtype=torch.float32) + elif (self.use_cute_dsl_bf16_gemm and is_sm_100f() + and self.weight.dtype == torch.bfloat16): input_2d = hidden_states.view(-1, hidden_states.shape[-1]) m, k = input_2d.shape n = self.weight.shape[0] @@ -930,6 +947,7 @@ def load_weights(self, "DeepseekV3Gate expects 'weight' and 'e_score_correction_bias' " "when partial loading is disabled") if w is not None: + self._weights_transformed = False self.weight.copy_(w[:]) if bias is not None: self.e_score_correction_bias.copy_(bias[:].to( @@ -992,7 +1010,8 @@ def __init__(self, fuse_routing_kernel=True, apply_routing=False, moe_backend=model_config.moe_backend, - use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm) + use_cute_dsl_bf16_gemm=model_config.use_cute_dsl_bf16_gemm, + locality_domain_policy=model_config.locality_domain_policy) # For MIXED_PRECISION, resolve the per-expert quant config (e.g. W4A8_AWQ) # instead of using the ambiguous global MIXED_PRECISION config. # For other cases (e.g. nvfp4, unquantized MTP layers), use @@ -1029,10 +1048,15 @@ def __init__(self, shared_quant_config = self._get_shared_experts_quant_config( model_config, layer_idx) - shared_model_config = model_config - if shared_quant_config is not model_config.quant_config: - shared_model_config = copy.copy(model_config) - shared_model_config.quant_config = shared_quant_config + # Shared experts overlap routed MoE on another stream. Disable locality + # domains for every precision so both paths do not contend for the same + # global partition streams. Only thaw the private copy of ModelConfig. + shared_model_config = copy.copy(model_config) + shared_model_config._frozen = False + shared_model_config.quant_config = shared_quant_config + shared_model_config.locality_domain_policy = LocalityDomainPolicy( + enabled=False) + shared_model_config._frozen = model_config._frozen # For shared experts, use the block size implied by their quant config. block_size = 1 @@ -1395,6 +1419,7 @@ def __init__(self, reduce_output=has_mlp_tp, use_cute_dsl_blockscaling_mm=model_config. use_cute_dsl_blockscaling_mm, + enable_locality_domain_bf16_linear=True, ) # On NVFP4 models, construct the boundary/prologue norms as NVFP4 norms. diff --git a/tensorrt_llm/_torch/models/modeling_deepseekv4.py b/tensorrt_llm/_torch/models/modeling_deepseekv4.py index 0e1a7638a702..85e545b9d80a 100644 --- a/tensorrt_llm/_torch/models/modeling_deepseekv4.py +++ b/tensorrt_llm/_torch/models/modeling_deepseekv4.py @@ -937,7 +937,7 @@ def load_flat_hc_weights(module, names: List[str]) -> bool: attn_module = all_named_modules[parent_module_name] _, v_b_proj = split_kv_b_proj(module.weight.data, is_scale=False) - attn_module.v_b_proj = nn.Parameter(v_b_proj, requires_grad=False) + attn_module._bind_v_b_proj_weight(v_b_proj) attn_module.k_b_proj_trans.data.copy_( k_b_proj_trans.reshape(attn_module.k_b_proj_trans.shape) @@ -1410,6 +1410,7 @@ def __init__( o_lora_rank=config.o_lora_rank, mapping_with_cp=mapping_with_cp, reduce_output=reduce_output, + enable_locality_domain_bf16_linear=True, ) diff --git a/tensorrt_llm/_torch/models/modeling_glm5_next.py b/tensorrt_llm/_torch/models/modeling_glm5_next.py index fed15d3ca435..bdbaebbbfc99 100644 --- a/tensorrt_llm/_torch/models/modeling_glm5_next.py +++ b/tensorrt_llm/_torch/models/modeling_glm5_next.py @@ -23,7 +23,7 @@ import copy import re -from collections.abc import Sequence +from collections.abc import Iterable, Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Callable @@ -702,8 +702,10 @@ def materialize(t: Any) -> torch.Tensor: target.copy_(tensor.to(device=device, dtype=target.dtype)) filled.add(local) - def mark_linear(mod_path: str) -> None: - for param_name in ("weight", "weight_scale", "bias"): + def mark_linear(mod_path: str, labels: Iterable[str] = ()) -> None: + # ``labels`` are extra tensors the Linear's load_weights consumed + # (e.g. the router gate's e_score_correction_bias). + for param_name in ("weight", "weight_scale", "bias", *labels): # mod_path is "" when the owner *is* the Linear (the LMHead). name = f"{mod_path}.{param_name}" if mod_path else param_name if name in params: @@ -711,7 +713,7 @@ def mark_linear(mod_path: str) -> None: for mod_path, group in linear_groups.items(): named_modules[mod_path].load_weights([group]) - mark_linear(mod_path) + mark_linear(mod_path, group) for fused_path, parts in fused_groups.items(): dest_mod = named_modules[fused_path] groups = [parts[i] for i in sorted(parts)] diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 21905050ebf0..a8a6f948ecbd 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -372,6 +372,8 @@ def apply_layerwise_quant_config( quant_config_dict = model_config.quant_config_dict if quant_config_dict is not None: for name, module in named_modules: + if getattr(module, '_skip_layerwise_quant_config', False): + continue if isinstance(module, (MoE, VanillaMoE)): for n, q in quant_config_dict.items(): # all linear layers inside FusedMoE share the same quant config diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py index dc1b6f0d5160..d19109ffb1c1 100644 --- a/tensorrt_llm/_torch/modules/gated_mlp.py +++ b/tensorrt_llm/_torch/modules/gated_mlp.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + from collections.abc import Callable from typing import Optional, Union @@ -41,6 +44,7 @@ def __init__( swiglu_alpha: Optional[float] = None, swiglu_beta: Optional[float] = None, split_gate_up: bool = False, + enable_locality_domain_bf16_linear: bool = False, ): super().__init__() @@ -67,6 +71,7 @@ def __init__( swiglu_beta) if swiglu_beta is not None else None config = config or ModelConfig() + locality_domain_policy = getattr(config, "locality_domain_policy", None) use_cute_dsl_bf16_gemm = getattr(config, "use_cute_dsl_bf16_gemm", False) self.mapping = config.mapping @@ -134,8 +139,12 @@ def __init__( use_cute_dsl_nvfp4_swiglu_blackwell=( use_cute_dsl_blockscaling_mm and activation == F.silu and not bias - and (swiglu_limit is None or swiglu_limit == float("inf"))), + and (swiglu_limit is None or swiglu_limit == float("inf")) + and self._is_plain_swiglu()), use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + enable_locality_domain_bf16_linear=( + use_cute_dsl_bf16_gemm and enable_locality_domain_bf16_linear), + locality_domain_policy=locality_domain_policy, disable_deep_gemm=disable_deep_gemm, use_custom_cublas_mm=use_custom_cublas_mm, ) @@ -186,6 +195,9 @@ def __init__( force_dynamic_quantization=config.force_dynamic_quantization, use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm, use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + enable_locality_domain_bf16_linear=( + use_cute_dsl_bf16_gemm and enable_locality_domain_bf16_linear), + locality_domain_policy=locality_domain_policy, disable_deep_gemm=disable_deep_gemm, use_custom_cublas_mm=use_custom_cublas_mm, ) @@ -336,7 +348,7 @@ def post_load_weights(self) -> None: return self._maybe_share_gate_up_quantize = True - def _can_fuse_gate_up_swiglu(self): + def _can_fuse_gate_up_swiglu(self) -> bool: """Check if fused GEMM + SwiGLU path is available. The projection owns the capability predicate because weight loading diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index adf8ad4dadaa..53e0344f21e1 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -40,6 +40,13 @@ from ...models.modeling_utils import QuantConfig from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE) +from ..locality_domain.layout import (make_bf16_linear_output_layout, + make_nvfp4_linear_output_layout) +from ..locality_domain.policy import (DISABLED_PLAN, LinearPartitionPlan, + LocalityDomainExecutionPlanner, + LocalityDomainPolicy) +from ..locality_domain.runtime import LocalityDomainRuntime +from ..locality_domain_utils import _copy_to_new_cuda_allocation from ..utils import (Fp4QuantizedTensor, get_model_extra_attrs, is_torch_compiling, replace_parameter_and_save_metadata, unswizzle_sf) @@ -3779,6 +3786,8 @@ def __init__( override_tp_sharding: Optional[Union[tuple[int, int], Dict[str, tuple[int, int]]]] = None, + locality_domain_policy: Optional[LocalityDomainPolicy] = None, + enable_locality_domain_bf16_linear: bool = False, ): """ Args: @@ -3812,6 +3821,8 @@ def __init__( self.use_cute_dsl_nvfp4_swiglu_blackwell = \ use_cute_dsl_nvfp4_swiglu_blackwell self.disable_deep_gemm = disable_deep_gemm + self.enable_locality_domain_bf16_linear = enable_locality_domain_bf16_linear + self._enable_gemm_allreduce_fusion = enable_gemm_allreduce_fusion self.fused_weight_shard_indices_mapping = fused_weight_shard_indices_mapping # Store NVFP4 GEMM allowed backends configuration # Read from model_extra_attrs if not explicitly provided (allows config via llm_api_options) @@ -3821,9 +3832,9 @@ def __init__( nvfp4_allowed_backends = model_attrs.get( 'nvfp4_gemm_allowed_backends') # Default: exclude cutedsl for faster build time - self.nvfp4_allowed_backends = nvfp4_allowed_backends or [ - 'cutlass', 'cublaslt', 'cuda_core' - ] + if nvfp4_allowed_backends is None: + nvfp4_allowed_backends = ['cutlass', 'cublaslt', 'cuda_core'] + self.nvfp4_allowed_backends = list(nvfp4_allowed_backends) if self.tp_mode not in (TensorParallelMode.ROW, TensorParallelMode.COLUMN, None): @@ -3901,26 +3912,17 @@ def __init__( self.use_cute_dsl_bf16_gemm = use_cute_dsl_bf16_gemm self.lora = lora - mpi_enabled = not mpi_disabled() - dtype_supported = self.dtype in (torch.float16, torch.bfloat16) - in_features_aligned = self.in_features % 128 == 0 - out_features_aligned = self.out_features % 64 == 0 - tp_valid = self.tp_mode is not None and self.tp_mode == TensorParallelMode.ROW and self.tp_size > 1 - quant_valid = quant_config_has_nvfp4_activation_quantization( - self.quant_config) - - device_supported = get_sm_version() >= 100 - enable_gemm_allreduce_fusion_env = (os.environ.get( - "TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "0") == "1") - - self.use_fused_gemm_allreduce = all([ - self.reduce_output, mpi_enabled, dtype_supported, - in_features_aligned, out_features_aligned, tp_valid, quant_valid, - device_supported, enable_gemm_allreduce_fusion, - enable_gemm_allreduce_fusion_env - ]) - if self.use_fused_gemm_allreduce: - self.use_fused_gemm_allreduce = ipc_nvls_supported() + if locality_domain_policy is None: + locality_domain_policy = (get_model_extra_attrs() + or {}).get("locality_domain_policy") + self._locality_domain_policy = locality_domain_policy or LocalityDomainPolicy( + ) + self.partition_plan = DISABLED_PLAN + self._locality_domain_runtime = None + self._locality_domain_weight_shards = None + self._locality_domain_added_cutedsl_backend = False + self._locality_domain_planned_quant_config = None + self._refresh_fused_gemm_allreduce() self.enable_cuda_core = False if torch.cuda.is_available(): @@ -4110,6 +4112,17 @@ def maybe_convert_to_torch_tensor( return maybe_convert_to_torch_tensor(weight, tuple(slice_obj)) def create_weights(self): + if self._weights_created: + if (self._locality_domain_policy.enabled and self.quant_config + is not self._locality_domain_planned_quant_config): + raise RuntimeError( + "Linear quant_config changed after weights were created; " + "reset _weights_created and recreate the weight schema.") + return + self._replan_locality_domain() + self.create_weights_impl() + + def create_weights_impl(self): if self._weights_created: return @@ -4169,7 +4182,8 @@ def can_use_cute_dsl_nvfp4_swiglu_blackwell(self) -> bool: and self.use_cute_dsl_blockscaling_mm and IS_CUTLASS_DSL_AVAILABLE and self.has_nvfp4_activation_quantization - and get_sm_version() in (100, 103) and not self.has_bias) + and get_sm_version() in (100, 103) and not self.has_bias + and not self.partition_plan.enabled) @property def requires_unquantized_activation(self): @@ -4229,6 +4243,12 @@ def apply_linear(self, bias, lora_params: Optional[dict] | None = None, layer_idx: Optional[int] | None = None): + if self._locality_domain_weight_shards is not None: + if self.lora is not None and bool(lora_params): + return LoraLayer.forward_with_base( + lambda: self._run_linear_locality_domain(input, bias), + (self.lora, ), input, lora_params, layer_idx) + return self._run_linear_locality_domain(input, bias) if self.lora is not None and bool(lora_params): output = LoraLayer.forward_with_base( lambda: self.quant_method.apply(self, input, bias), @@ -4275,7 +4295,14 @@ def forward( layer_idx: Optional[int] = None, ) -> torch.Tensor: if self.tp_mode == TensorParallelMode.ROW: - use_fused_gemm_allreduce = self.use_fused_gemm_allreduce and lora_params is None + use_locality_domain = self.partition_plan.enabled + # The fused GEMM-allreduce and NCCL-window GEMM paths consume the + # module's full weight tensors. locality domain releases those tensors after + # creating localized shards, so ROW-TP must first run the + # shard-aware GEMM and then all-reduce its partial output. + use_fused_gemm_allreduce = (not use_locality_domain + and self.use_fused_gemm_allreduce + and lora_params is None) if use_fused_gemm_allreduce and all_reduce_params is not None: use_fused_gemm_allreduce = all_reduce_params.enable_allreduce and all_reduce_params.fusion_op == AllReduceFusionOp.NONE @@ -4294,7 +4321,8 @@ def forward( # (ClassVar); a failed window allocation falls back gracefully # inside the C++ allocate_output. use_nccl_symmetric_memory_window = ( - self.all_reduce is not None and self.quant_method. + not use_locality_domain and self.all_reduce is not None + and self.quant_method. supports_nccl_symmetric_memory_window_output and self.all_reduce.uses_nccl_symmetric_memory_window() and not (self.lora is not None and lora_params)) @@ -4345,6 +4373,22 @@ def transform_weights(self) -> None: if self._weights_transformed: return self.quant_method.transform_weights(self) + # Localized allocations and full-weight release belong to the mutating + # stage so both staged ModelLoader walks and post_load_weights run them. + if self._locality_domain_runtime is not None: + shards = self._split_weights_for_locality_domain() + self._locality_domain_runtime.prepare_for_capture( + self.partition_plan) + self._locality_domain_weight_shards = shards + self._save_locality_domain_reload_metadata() + # Free full weights — shards on localized memory are the source of + # truth now. Keep bias because forward() uses it to preserve the + # module's logical bias contract and to support fused allreduce bias. + self.weight = Parameter(self.weight.new_empty(0), + requires_grad=False) + if self.partition_plan.op_kind == "nvfp4_linear": + self.weight_scale = Parameter(self.weight_scale.new_empty(0), + requires_grad=False) self._weights_transformed = True def cache_derived_state(self) -> None: @@ -4358,7 +4402,339 @@ def pre_reload_weights(self): assert hasattr( self.quant_method, "pre_reload_weights" ), "pre_reload_weights is not supported for this quant method" + had_locality_domain_shards = self._locality_domain_weight_shards is not None + if had_locality_domain_shards: + # Drop localized sources before allocating the full reload schema. + # post_load_weights() will rebuild the localized shards after the + # new weights have been loaded and transformed. + self._locality_domain_weight_shards = None self.quant_method.pre_reload_weights(self) + if had_locality_domain_shards: + # transform_weights() must treat the new full parameters as the + # first generation. In particular, padding transforms populate + # fresh reusable metadata instead of referencing released tensors. + self.rebuild_tensor_metadata = {} + self._weights_transformed = False + + def _refresh_fused_gemm_allreduce(self) -> None: + """Recompute the fused GEMM/all-reduce gate from final quantization.""" + quant_valid = quant_config_has_nvfp4_activation_quantization( + self.quant_config) + enabled = all([ + self.reduce_output, + not mpi_disabled(), + self.dtype in (torch.float16, torch.bfloat16), + self.in_features % 128 == 0, + self.out_features % 64 == 0, + self.tp_mode == TensorParallelMode.ROW, + self.tp_size > 1, + quant_valid, + get_sm_version() >= 100, + self._enable_gemm_allreduce_fusion, + os.environ.get("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "0") == "1", + ]) + self.use_fused_gemm_allreduce = enabled and ipc_nvls_supported() + + def _replan_locality_domain(self) -> None: + """Plan locality domain execution from the final effective Linear config.""" + if self._weights_created: + raise RuntimeError( + "Cannot replan locality domain Linear while its weight schema is live; " + "reset and recreate the module weights after changing " + "quantization.") + if self._locality_domain_weight_shards is not None: + raise RuntimeError( + "Cannot replan locality domain Linear after localized weight shards were " + "created; reload the module weights before changing its " + "quantization or locality domain policy.") + + planner = LocalityDomainExecutionPlanner(self._locality_domain_policy) + plan = planner.plan_linear( + self.in_features, + self.out_features, + self.quant_config, + self.weights_loading_config.weight_mode, + dtype=self.dtype, + use_cute_dsl_bf16_gemm=self.use_cute_dsl_bf16_gemm, + enable_locality_domain_bf16_linear=self. + enable_locality_domain_bf16_linear, + ) + if (plan.enabled and self.weights_loading_config.weight_mode + == WeightMode.FUSED_GATE_UP_LINEAR): + half_n = self.out_features // 2 + expected_mapping = { + "gate": (0, half_n), + "up": (half_n, half_n), + } + if self.fused_weight_shard_indices_mapping != expected_mapping: + plan = LinearPartitionPlan( + enabled=False, + reason_if_disabled= + ("locality domain fused gate/up Linear requires canonical " + "[gate | up] shard mapping"), + ) + + if self._locality_domain_added_cutedsl_backend: + self.nvfp4_allowed_backends.remove("cutedsl") + self._locality_domain_added_cutedsl_backend = False + if (plan.enabled and plan.op_kind == "nvfp4_linear" + and "cutedsl" not in self.nvfp4_allowed_backends): + self.nvfp4_allowed_backends.append("cutedsl") + self._locality_domain_added_cutedsl_backend = True + + self.partition_plan = plan + self._locality_domain_runtime = (LocalityDomainRuntime( + plan.num_partitions) if plan.enabled else None) + self._locality_domain_planned_quant_config = self.quant_config + self._refresh_fused_gemm_allreduce() + + def _run_linear_locality_domain(self, input, bias): + if self.partition_plan.op_kind == "nvfp4_linear": + return self._run_nvfp4_linear_locality_domain(input, bias) + if self.partition_plan.op_kind == "bf16_linear": + return self._run_bf16_linear_locality_domain(input, bias) + raise RuntimeError( + f"Unsupported locality domain Linear operation: {self.partition_plan.op_kind}" + ) + + def _run_nvfp4_linear_locality_domain(self, input, bias): + """Execute a partitioned linear across both locality domain partitions. + + The composite CuTe DSL op tunes with both partitions active, then + launches the same tactic on both partition streams. Each kernel writes + directly into its slice of a shared, strided output buffer. + + Args: + bias: Caller-controlled bias. None means skip bias (e.g. ROW + parallel tp_rank > 0, or fused into allreduce). + """ + shards = self._locality_domain_weight_shards + layout = self.partition_plan.layout + assert layout is not None, "locality domain Linear requires partition layout metadata" + num_p = layout.num_partitions + assert num_p == 2, "locality domain Linear requires exactly two partitions" + n = layout.logical_axis_extent + + # Handle multi-dimensional inputs (e.g., 3D: batch, seq, hidden) + original_shape = None + if isinstance(input, Fp4QuantizedTensor) and input.fp4_tensor.dim() > 2: + original_shape = input.fp4_tensor.shape + input = Fp4QuantizedTensor( + fp4_tensor=input.fp4_tensor.reshape(-1, + input.fp4_tensor.shape[-1]), + scaling_factor=input.scaling_factor, + is_sf_swizzled=input.is_sf_swizzled, + unquantized_hidden_states=input.unquantized_hidden_states, + ) + elif isinstance(input, tuple) and input[0].dim() > 2: + original_shape = input[0].shape + input = ( + input[0].reshape(-1, input[0].shape[-1]), + input[1], + ) + elif not isinstance(input, + (tuple, Fp4QuantizedTensor)) and input.dim() > 2: + original_shape = input.shape + input = input.reshape(-1, input.shape[-1]) + + # Quantize input once (shared across partitions) + act_fp4, act_sf, alpha = self.quant_method._input_prepare(self, input) + m = act_fp4.shape[0] + + # Each shard's N may be padded beyond out_features / num_partitions. + shard_n = layout.per_partition_axis_extent(padded=True) + assert shards[0]['weight'].size(0) == shard_n + full_n = layout.padded_axis_extent + + # Output buffer: full_n to accommodate padding, truncated at the end + output = torch.empty(m, full_n, dtype=self.dtype, device=act_fp4.device) + + torch.ops.trtllm.cute_dsl_nvfp4_gemm_locality_domain_inplace_rubin( + act_fp4, + shards[0]['weight'], + shards[1]['weight'], + act_sf, + shards[0]['weight_scale'], + shards[1]['weight_scale'], + alpha, + self.dtype, + False, + True, + output, + ) + + # Truncate padded columns if needed + if full_n > n: + output = output[:, :n].contiguous() + + if bias is not None: + output.add_(bias) + + if original_shape is not None: + output = output.reshape(*original_shape[:-1], output.shape[-1]) + + return output + + def _run_bf16_linear_locality_domain( + self, + input: torch.Tensor, + bias: Optional[torch.Tensor], + *, + output_dtype: torch.dtype = torch.bfloat16, + ) -> torch.Tensor: + """Execute a BF16 Linear with output-N split across two locality domains.""" + shards = self._locality_domain_weight_shards + layout = self.partition_plan.layout + assert shards is not None + assert layout is not None, "locality domain BF16 Linear requires layout metadata" + assert layout.num_partitions == 2 + assert input.dtype == torch.bfloat16 + + original_shape = input.shape + input_2d = input.reshape(-1, input.shape[-1]).contiguous() + output = torch.empty( + input_2d.shape[0], + layout.logical_axis_extent, + dtype=output_dtype, + device=input.device, + ) + if input_2d.shape[0] > 0: + torch.ops.trtllm.cute_dsl_bf16_gemm_locality_domain_inplace_rubin( + input_2d, + shards[0]["weight"], + shards[1]["weight"], + output, + True, + ) + if bias is not None: + output.add_(bias) + return output.reshape(*original_shape[:-1], output.shape[-1]) + + def _save_locality_domain_reload_metadata(self) -> None: + """Keep only the original schemas needed to rebuild released weights.""" + parameter_names = ["weight"] + if self.partition_plan.op_kind == "nvfp4_linear": + parameter_names.append("weight_scale") + + for parameter_name in parameter_names: + metadata = self.rebuild_tensor_metadata.get(parameter_name) + if metadata is None: + meta_tensor = getattr(self, parameter_name).to("meta") + else: + meta_tensor = metadata["meta"] + # Do not retain metadata["param"]: it aliases the transformed full + # weight and would defeat releasing it after localization. + self.rebuild_tensor_metadata[parameter_name] = {"meta": meta_tensor} + + def _split_weights_for_locality_domain(self): + if self.partition_plan.op_kind == "nvfp4_linear": + return self._split_nvfp4_weights_for_locality_domain() + if self.partition_plan.op_kind == "bf16_linear": + return self._split_bf16_weights_for_locality_domain() + raise RuntimeError( + f"Unsupported locality domain Linear operation: {self.partition_plan.op_kind}" + ) + + def _split_nvfp4_weights_for_locality_domain(self): + """Split full-N weights into per-partition shards on localized memory. + + Called after normal load + post_load so weight/weight_scale are final + (including padding and interleaving). + + Weight (dim=0 is N, may be padded): split along dim=0. + Weight_scale: unswizzle to 2D → split unpadded rows → re-swizzle each + partition (swizzle_sf handles per-partition padding internally). + """ + num_p = self.partition_plan.num_partitions + layout = make_nvfp4_linear_output_layout( + self.out_features, + self.in_features, + num_p, + padded_out_features=self.weight.size(0), + ) + + # Validate padding split invariant: splitting the full weight evenly + # must not introduce per-partition padding gaps in the output layout. + layout_reason = layout.disabled_reason_for_padding_free_split() + assert layout_reason is None, layout_reason + assert self.partition_plan.layout == layout, ( + f"Runtime layout {layout} does not match planner layout " + f"{self.partition_plan.layout}") + if (self.weights_loading_config.weight_mode == + WeightMode.FUSED_GATE_UP_LINEAR): + half_n = self.out_features // 2 + assert self.fused_weight_shard_indices_mapping == { + "gate": (0, half_n), + "up": (half_n, half_n), + }, ("locality domain fused gate/up Linear requires canonical [gate | up] " + "row layout") + + # Unswizzle weight_scale to 2D, split unpadded N rows, re-swizzle. + # We split on unpadded rows because the swizzle pattern is not simply + # splittable; re-swizzle handles per-partition padding internally. + sv = self.scaling_vector_size + scale_rows = fp4_utils.pad_up(self.out_features, 128) + ws_unswizzled = unswizzle_sf( + self.weight_scale.data, + scale_rows, + self.in_features, + sv, + ) + + shards = [] + for pid in range(num_p): + padded_slice = layout.partition_axis_slice(pid, padded=True) + logical_slice = layout.partition_axis_slice(pid, padded=False) + with self._locality_domain_runtime.partition_weight_context(pid): + ws_part = ws_unswizzled[logical_slice].contiguous() + # The interleave op creates its output while the partition's + # localized memory pool is active. + ws_part_swizzled = torch.ops.trtllm.block_scale_interleave( + ws_part.unsqueeze(0)) + shard = { + 'weight': + _copy_to_new_cuda_allocation(self.weight[padded_slice]), + 'weight_scale': + ws_part_swizzled, + } + shards.append(shard) + return shards + + def _split_bf16_weights_for_locality_domain(self): + """Split BF16 output rows into fresh allocations in localized pools.""" + assert self._locality_domain_runtime is not None + assert self.weight.dtype == torch.bfloat16 + layout = make_bf16_linear_output_layout( + self.out_features, + self.in_features, + self.partition_plan.num_partitions, + ) + layout_reason = layout.disabled_reason_for_padding_free_split() + assert layout_reason is None, layout_reason + assert self.partition_plan.layout == layout, ( + f"Runtime layout {layout} does not match planner layout " + f"{self.partition_plan.layout}") + assert tuple(self.weight.shape) == layout.logical_shape + if (self.weights_loading_config.weight_mode == + WeightMode.FUSED_GATE_UP_LINEAR): + half_n = self.out_features // 2 + assert self.fused_weight_shard_indices_mapping == { + "gate": (0, half_n), + "up": (half_n, half_n), + }, ("locality domain BF16 fused gate/up Linear requires canonical " + "[gate | up] row layout") + + shards = [] + for partition_id in range(layout.num_partitions): + weight_slice = layout.partition_axis_slice(partition_id, + padded=False) + with self._locality_domain_runtime.partition_weight_context( + partition_id): + shards.append({ + "weight": + _copy_to_new_cuda_allocation(self.weight[weight_slice]) + }) + return shards def is_static_nvfp4_input_eligible(linear) -> bool: diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index e50f7a362312..efa2506b5c5b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py @@ -460,24 +460,24 @@ def forward(self, tile_size = tactic else: tile_size = 128 - recv_expert_count = None - forward_inputs = inputs - if self.use_direct_expert_metadata: - recv_expert_count = inputs[-1] - forward_inputs = inputs[:-1] - num_rows = forward_inputs[0].size(0) - if num_rows % self.num_local_experts != 0: - raise ValueError( - "Expert-major input rows must be divisible by the number " - "of local experts") - deep_ep_expert_capacity = num_rows // self.num_local_experts - else: - deep_ep_expert_capacity = None + if not self.use_direct_expert_metadata: + return self.forward_impl(*inputs, + enable_alltoall=self.enable_alltoall, + tile_size=tile_size) + + recv_expert_count = inputs[-1] + forward_inputs = inputs[:-1] + num_rows = forward_inputs[0].size(0) + if num_rows % self.num_local_experts != 0: + raise ValueError( + "Expert-major input rows must be divisible by the number " + "of local experts") + deep_ep_expert_capacity = num_rows // self.num_local_experts # The locality-domain impl does not take the count-native arguments. count_native_kwargs = {} if self.use_locality_domain else dict( recv_expert_count=recv_expert_count, deep_ep_expert_capacity=deep_ep_expert_capacity, - use_count_native_expert_metadata=self.use_direct_expert_metadata) + use_count_native_expert_metadata=True) return self.forward_impl(*forward_inputs, enable_alltoall=self.enable_alltoall, tile_size=tile_size, @@ -2036,6 +2036,8 @@ def run_moe( def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): + if self._locality_domain_weight_shards is not None: + self.pre_reload_weights() super().load_weights(weights, allow_partial_loading=allow_partial_loading) # Keep DWDP registration after base weight loading. This preserves @@ -2051,6 +2053,8 @@ def transform_weights(self) -> None: super().transform_weights() # Split full weights into per-partition halves on localized memory if self._locality_domain_runtime is not None: + # Refresh views of finalized scales before copying their shards. + self.quant_method.setup_quant_scales(self) self._locality_domain_weight_shards = self._split_weights_for_locality_domain( ) self._release_full_weights_after_locality_domain_split() @@ -2069,6 +2073,13 @@ def cache_derived_state(self) -> None: # Resolve the borrowed remainder stream now, never during capture. self._get_reserved_moe_output_memset_stream() + def pre_reload_weights(self) -> None: + # The quant method rebuilds the original full-weight schemas. Drop + # shards so the next transform uses the newly loaded checkpoint. + super().pre_reload_weights() + self._locality_domain_weight_shards = None + self._weights_transformed = False + def _release_full_weights_after_locality_domain_split(self): """Release full tensors that are replaced by localized locality domain shards.""" for param_name in ( @@ -2080,6 +2091,11 @@ def _release_full_weights_after_locality_domain_split(self): param = getattr(self, param_name, None) if param is None: continue + metadata = self.rebuild_tensor_metadata.get(param_name) + meta_tensor = param.to( + "meta") if metadata is None else metadata["meta"] + # Retain the reload schema without a reference to the full weight. + self.rebuild_tensor_metadata[param_name] = {"meta": meta_tensor} setattr( self, param_name, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 13e5a6c1ec87..1fda3330ea7d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -42,6 +42,7 @@ from tensorrt_llm.quantization.utils.fp4_utils import float4_e2m1x2 from ...llmapi.llm_args import LoadFormat +from ..locality_domain.policy import LocalityDomainPolicy from ..model_config import ModelConfig from ..models import AutoModelForCausalLM from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader @@ -934,6 +935,14 @@ def load( post_transform_config_identity = PostTransformConfigIdentity.from_model_config( config) load_format = self.llm_args.load_format + locality_domain_policy = getattr(config, "locality_domain_policy", None) + if (load_format == LoadFormat.GMS + and getattr(locality_domain_policy, "enabled", False)): + raise ValueError( + "LoadFormat.GMS is incompatible with locality domain localized weights. " + "GMS shares registered parameters, while locality domain execution " + "releases those parameters after creating partition-local " + "weight shards.") self._startup_metadata = _checkpoint_startup_metadata( checkpoint_loader, load_format) @@ -1602,6 +1611,11 @@ def _qualify_post_transform_profile( enabled_features = set() if loads_draft_weights: enabled_features.add(PostTransformFeature.SEPARATE_DRAFT_MODEL) + locality_domain_policy = getattr(model.model_config, + "locality_domain_policy", None) + if getattr(locality_domain_policy, "enabled", False): + enabled_features.add( + PostTransformFeature.LOCALITY_DOMAIN_LOCALIZED_WEIGHTS) return cls._post_transform_profile_registry().qualify( root_model_class=type(model), architecture=config_identity.architecture, @@ -1922,6 +1936,8 @@ def _load_and_validate_config( multimodal_config=self.llm_args.multimodal_config, use_cute_dsl_bf16_bmm=self.llm_args.use_cute_dsl_bf16_bmm, use_cute_dsl_bf16_gemm=self.llm_args.use_cute_dsl_bf16_gemm, + locality_domain_policy=LocalityDomainPolicy( + enabled=self.llm_args.enable_locality_domains), ) # Only pass model_kwargs if it's explicitly set (not None) diff --git a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py index 3314707c47fd..1bd14e39d76c 100644 --- a/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py +++ b/tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py @@ -351,6 +351,7 @@ class PostTransformFeature(str, Enum): """Optional lifecycle features that require explicit qualification.""" SEPARATE_DRAFT_MODEL = "separate_draft_model" + LOCALITY_DOMAIN_LOCALIZED_WEIGHTS = "localized_weights" class PostTransformQualificationReason(str, Enum): diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 20875cb003ad..52c668736e2d 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -289,6 +289,7 @@ def get_llm_args( agent_percentage: float = 0.0, agent_types: Optional[str] = None, explicit_cli_keys: Optional[Set[str]] = None, + enable_locality_domains: bool = False, **llm_args_extra_dict: Any): explicit_cli_keys = explicit_cli_keys or set() @@ -358,6 +359,8 @@ def get_llm_args( enable_chunked_prefill, "enable_attention_dp": enable_attention_dp, + "enable_locality_domains": + enable_locality_domains, "revision": revision, "reasoning_parser": @@ -1254,6 +1257,12 @@ def launch_visual_gen_server( default=False, help="Enable attention data parallel.", status="beta") +@stability_option( + "--enable_locality_domains", + is_flag=True, + default=False, + help="Enable locality domain execution for supported PyTorch backend ops.", + status="prototype") @stability_option("--media_io_kwargs", type=str, default=None, @@ -1401,6 +1410,7 @@ def serve( served_model_name: Optional[str], visual_gen_args: Optional[str], report_addr: Optional[str], + enable_locality_domains: bool = False, ) -> None: """Running an OpenAI API compatible server @@ -1517,6 +1527,7 @@ def _serve_llm(): otlp_traces_endpoint=otlp_traces_endpoint, enable_chunked_prefill=enable_chunked_prefill, enable_attention_dp=enable_attention_dp, + enable_locality_domains=enable_locality_domains, video_pruning_rate=video_pruning_rate, telemetry=telemetry, agent_percentage=agent_percentage, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f241528a2b31..f3c3bd9c0043 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6365,6 +6365,12 @@ def is_partial_model_loading(self) -> bool: description="If true, use CuTe DSL fp8 blockscaling bmm implementation.", status="prototype") # bf16 cute dsl configs + enable_locality_domains: bool = Field( + default=False, + description= + "Enable locality domain execution for supported PyTorch backend ops.", + status="prototype") + use_cute_dsl_bf16_bmm: bool = Field( default=False, description= @@ -7096,6 +7102,16 @@ def validate_gms_config(self) -> 'TorchLlmArgs': "enable GPU Memory Service.", self.load_format.name) return self + @model_validator(mode="after") + def validate_gms_locality_domain_compat(self) -> 'TorchLlmArgs': + """Reject localized weights that cannot be shared through the GMS pool.""" + if self.load_format == LoadFormat.GMS and self.enable_locality_domains: + raise ValueError( + "LoadFormat.GMS is incompatible with enable_locality_domains=True. " + "Localized weights are allocated outside the shared GMS pool. " + "Disable enable_locality_domains or use LoadFormat.AUTO.") + return self + @model_validator(mode="after") def validate_gms_moe_compat(self) -> 'TorchLlmArgs': """Reject ``LoadFormat.GMS`` combined with a MoE load balancer. diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index c836fff6843c..340b81bed937 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -365,6 +365,11 @@ "kind": "value", "path": "enable_lm_head_tp_in_adp" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "enable_locality_domains" + }, { "capture_policy": "bool", "kind": "value", diff --git a/tests/unittest/_torch/attention/test_mla_locality_domain.py b/tests/unittest/_torch/attention/test_mla_locality_domain.py new file mode 100644 index 000000000000..5fefc8ff40e4 --- /dev/null +++ b/tests/unittest/_torch/attention/test_mla_locality_domain.py @@ -0,0 +1,417 @@ +# 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. + +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.attention import mla as mla_module +from tensorrt_llm._torch.attention.backends.interface import PositionalEmbeddingParams, RopeParams +from tensorrt_llm._torch.attention.mla import MLA +from tensorrt_llm._torch.locality_domain.policy import ( + LocalityDomainExecutionPlanner, + LocalityDomainPolicy, +) +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.functional import PositionEmbeddingType + + +class _FakeAttention(nn.Module): + def support_fused_rope(self) -> bool: + return True + + +class _RecordingLocalityDomainRuntime: + def __init__(self) -> None: + self.num_partitions = 2 + self.entered_partitions: list[int] = [] + self.exited_partitions: list[int] = [] + self.prepared_plans: list[object] = [] + + @contextmanager + def partition_weight_context(self, partition_id: int) -> Iterator[None]: + self.entered_partitions.append(partition_id) + yield + self.exited_partitions.append(partition_id) + + def prepare_for_capture(self, plan: object) -> None: + self.prepared_plans.append(plan) + + +def _make_bare_absorption_mla( + *, + policy_enabled: bool = True, + use_cute_dsl_bf16_bmm: bool = True, + is_deepseek_v4: bool = False, + weight_dtype: torch.dtype = torch.bfloat16, + qk_nope_head_dim: int = 128, +) -> tuple[MLA, _RecordingLocalityDomainRuntime]: + mla = MLA.__new__(MLA) + nn.Module.__init__(mla) + mla.is_deepseek_v4 = is_deepseek_v4 + mla.use_cute_dsl_bf16_bmm = use_cute_dsl_bf16_bmm + mla._locality_domain_policy = LocalityDomainPolicy( + enabled=policy_enabled, + allowed_ops=frozenset({"bf16_bmm"}), + ) + runtime = _RecordingLocalityDomainRuntime() + mla._locality_domain_runtime = runtime + mla._locality_domain_k_b_proj_trans_shards = None + mla._locality_domain_v_b_proj_shards = None + mla._weights_transformed = False + mla.kv_b_proj = SimpleNamespace( + quant_config=None, partition_plan=SimpleNamespace(enabled=False) + ) + mla.num_heads_tp = 2 + mla.num_heads_tp_cp = 2 + mla.qk_nope_head_dim = qk_nope_head_dim + mla.kv_lora_rank = 512 + mla.v_head_dim = 128 + + k_values = torch.arange( + mla.num_heads_tp * mla.kv_lora_rank * mla.qk_nope_head_dim, + dtype=torch.float32, + ).view(mla.num_heads_tp, mla.kv_lora_rank, mla.qk_nope_head_dim) + v_values = torch.arange( + mla.num_heads_tp_cp * mla.v_head_dim * mla.kv_lora_rank, + dtype=torch.float32, + ).view(mla.num_heads_tp_cp, mla.v_head_dim, mla.kv_lora_rank) + mla.k_b_proj_trans = nn.Parameter(k_values.to(weight_dtype), requires_grad=False) + mla.v_b_proj = nn.Parameter(v_values.to(weight_dtype), requires_grad=False) + return mla, runtime + + +def _make_mla_with_config( + config: ModelConfig, + q_lora_rank: int | None, + *, + enable_locality_domain_bf16_linear: bool = True, +) -> MLA: + position_embedding = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=RopeParams(dim=2, max_positions=8), + ) + return MLA( + hidden_size=8, + num_attention_heads=2, + num_key_value_heads=1, + qk_nope_head_dim=2, + qk_rope_head_dim=2, + v_head_dim=2, + q_lora_rank=q_lora_rank, + kv_lora_rank=4, + predicted_tokens_per_seq=1, + max_position_embeddings=8, + bias=False, + pos_embd_params=position_embedding, + layer_idx=0, + dtype=torch.bfloat16, + config=config, + o_lora_rank=2, + enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear, + ) + + +@pytest.mark.parametrize( + ("q_lora_rank", "enable_locality_domain_bf16_linear"), + [ + pytest.param(4, True, id="q-lora"), + pytest.param(None, True, id="fused-q-projection"), + pytest.param(4, False, id="explicit-opt-out"), + ], +) +def test_mla_wires_locality_domain_policy_to_bf16_linear_projections( + q_lora_rank: int | None, + enable_locality_domain_bf16_linear: bool, +) -> None: + allowed_ops = {"bf16_linear"} + if enable_locality_domain_bf16_linear: + allowed_ops.add("bf16_bmm") + policy = LocalityDomainPolicy(enabled=True, allowed_ops=frozenset(allowed_ops)) + config = ModelConfig( + skip_create_weights_in_init=True, + use_cute_dsl_bf16_gemm=True, + locality_domain_policy=policy, + ) + + with patch.object(mla_module, "create_attention", return_value=_FakeAttention()): + mla = _make_mla_with_config( + config, + q_lora_rank, + enable_locality_domain_bf16_linear=enable_locality_domain_bf16_linear, + ) + + assert mla._locality_domain_policy is policy + for projection in (mla.q_b_proj, mla.kv_b_proj, mla.o_proj): + assert projection.use_cute_dsl_bf16_gemm + assert projection.enable_locality_domain_bf16_linear is enable_locality_domain_bf16_linear + assert projection._locality_domain_policy is policy + assert not mla.kv_a_proj_with_mqa.enable_locality_domain_bf16_linear + if q_lora_rank is None: + assert mla.q_b_proj is mla.q_proj + + +@pytest.mark.parametrize("partition_enabled", [True, False]) +def test_bind_v_b_proj_weight_respects_bf16_linear_partition_plan( + partition_enabled: bool, +) -> None: + mla = MLA.__new__(MLA) + nn.Module.__init__(mla) + mla.kv_b_proj = SimpleNamespace( + partition_plan=SimpleNamespace( + enabled=partition_enabled, + op_kind="bf16_linear" if partition_enabled else None, + ) + ) + source = torch.arange(24, dtype=torch.bfloat16).view(4, 6) + weight = source[1:] + + mla._bind_v_b_proj_weight(weight) + + torch.testing.assert_close(mla.v_b_proj, weight) + aliases_weight = ( + mla.v_b_proj.untyped_storage().data_ptr() == weight.untyped_storage().data_ptr() + ) + assert aliases_weight is not partition_enabled + + +def test_transform_builds_fresh_localized_absorption_shards( + monkeypatch: pytest.MonkeyPatch, +) -> None: + plan_calls = [] + + def _plan_bf16_bmm( + _planner: LocalityDomainExecutionPlanner, + *, + dtype: torch.dtype, + use_cute_dsl_bf16_bmm: bool, + ) -> SimpleNamespace: + plan_calls.append((dtype, use_cute_dsl_bf16_bmm)) + return SimpleNamespace(enabled=True, num_partitions=2) + + monkeypatch.setattr( + mla_module.LocalityDomainExecutionPlanner, + "plan_bf16_bmm", + _plan_bf16_bmm, + ) + mla, runtime = _make_bare_absorption_mla() + original_k = mla.k_b_proj_trans + original_v = mla.v_b_proj + + mla.transform_weights() + + k_shards = mla._locality_domain_k_b_proj_trans_shards + v_shards = mla._locality_domain_v_b_proj_shards + assert k_shards is not None + assert v_shards is not None + assert [tuple(shard.shape) for shard in k_shards] == [ + (mla.num_heads_tp, 256, mla.qk_nope_head_dim), + (mla.num_heads_tp, 256, mla.qk_nope_head_dim), + ] + assert [tuple(shard.shape) for shard in v_shards] == [ + (mla.num_heads_tp_cp, 64, mla.kv_lora_rank), + (mla.num_heads_tp_cp, 64, mla.kv_lora_rank), + ] + torch.testing.assert_close(k_shards[0], original_k[:, :256]) + torch.testing.assert_close(k_shards[1], original_k[:, 256:]) + torch.testing.assert_close(v_shards[0], original_v[:, :64]) + torch.testing.assert_close(v_shards[1], original_v[:, 64:]) + assert all(shard.is_contiguous() for shard in (*k_shards, *v_shards)) + assert all( + shard.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + for shard, source in ( + (k_shards[0], original_k), + (k_shards[1], original_k), + (v_shards[0], original_v), + (v_shards[1], original_v), + ) + ) + assert runtime.entered_partitions == [0, 1, 0, 1] + assert runtime.exited_partitions == [0, 1, 0, 1] + assert len(runtime.prepared_plans) == 1 + assert plan_calls == [(torch.bfloat16, True)] + assert mla.k_b_proj_trans is original_k + assert mla.v_b_proj is original_v + + # Repeated staged/legacy hooks must preserve capture-stable shard storage. + mla.cache_derived_state() + mla.post_load_weights() + assert mla._locality_domain_k_b_proj_trans_shards is k_shards + assert mla._locality_domain_v_b_proj_shards is v_shards + assert len(runtime.prepared_plans) == 1 + + old_k_shards = k_shards + with torch.no_grad(): + original_k.fill_(3) + original_v.fill_(5) + # The checkpoint loader binds the new V weight after loading kv_b_proj. + # It must invalidate the derived shards without a manual flag reset. + mla._bind_v_b_proj_weight(original_v) + mla.post_load_weights() + + rebuilt_k_shards = mla._locality_domain_k_b_proj_trans_shards + rebuilt_v_shards = mla._locality_domain_v_b_proj_shards + assert rebuilt_k_shards is not None + assert rebuilt_v_shards is not None + assert ( + rebuilt_k_shards[0].untyped_storage().data_ptr() + != old_k_shards[0].untyped_storage().data_ptr() + ) + torch.testing.assert_close(rebuilt_k_shards[0], original_k[:, :256]) + torch.testing.assert_close(rebuilt_v_shards[1], original_v[:, 64:]) + assert plan_calls == [(torch.bfloat16, True), (torch.bfloat16, True)] + assert len(runtime.prepared_plans) == 2 + + +@pytest.mark.parametrize( + ("is_deepseek_v4", "qk_nope_head_dim"), + [ + pytest.param(False, 127, id="unaligned-qk-nope-head-dim"), + pytest.param(True, 128, id="deepseek-v4"), + ], +) +def test_absorption_shards_reject_unsupported_weights_before_planning( + monkeypatch: pytest.MonkeyPatch, + is_deepseek_v4: bool, + qk_nope_head_dim: int, +) -> None: + def _unexpected_plan(*_args: object, **_kwargs: object) -> None: + pytest.fail("unsupported absorption weights must be rejected before planning") + + monkeypatch.setattr( + mla_module.LocalityDomainExecutionPlanner, + "plan_bf16_bmm", + _unexpected_plan, + ) + mla, runtime = _make_bare_absorption_mla( + is_deepseek_v4=is_deepseek_v4, + qk_nope_head_dim=qk_nope_head_dim, + ) + + if is_deepseek_v4: + # DSV4 sparse hooks do not create absorption projection weights. + del mla.k_b_proj_trans + del mla.v_b_proj + mla.transform_weights() + + assert mla._locality_domain_k_b_proj_trans_shards is None + assert mla._locality_domain_v_b_proj_shards is None + assert runtime.entered_partitions == [] + + +def test_absorption_shards_forward_inputs_to_disabled_plan( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def _plan_bf16_bmm( + planner: LocalityDomainExecutionPlanner, + *, + dtype: torch.dtype, + use_cute_dsl_bf16_bmm: bool, + ) -> SimpleNamespace: + assert not planner.policy.enabled + assert dtype == torch.float8_e4m3fn + assert not use_cute_dsl_bf16_bmm + return SimpleNamespace(enabled=False) + + monkeypatch.setattr( + mla_module.LocalityDomainExecutionPlanner, + "plan_bf16_bmm", + _plan_bf16_bmm, + ) + mla, runtime = _make_bare_absorption_mla( + policy_enabled=False, + use_cute_dsl_bf16_bmm=False, + weight_dtype=torch.float8_e4m3fn, + ) + + mla.transform_weights() + + assert mla._locality_domain_k_b_proj_trans_shards is None + assert mla._locality_domain_v_b_proj_shards is None + assert runtime.entered_partitions == [] + + +@pytest.mark.parametrize( + ("use_cute", "sm_version", "with_shards", "expected_op"), + [ + (True, 107, True, "locality_domain"), + (True, 107, False, "rubin"), + (True, 100, True, "blackwell"), + (True, 103, True, "blackwell"), + (True, 90, True, "bmm_out"), + (False, 107, True, "bmm_out"), + ], +) +def test_bf16_bmm_dispatch_and_fallback( + monkeypatch: pytest.MonkeyPatch, + use_cute: bool, + sm_version: int, + with_shards: bool, + expected_op: str, +) -> None: + mla = MLA.__new__(MLA) + nn.Module.__init__(mla) + mla.use_cute_dsl_bf16_bmm = use_cute + monkeypatch.setattr(mla_module, "get_sm_version", lambda: sm_version) + monkeypatch.setattr(mla_module, "is_sm_100f", lambda: 100 <= sm_version < 110) + + calls = [] + + def _record(name: str) -> Callable[..., None]: + def _op(*args: object) -> None: + calls.append((name, args)) + + return _op + + monkeypatch.setattr( + torch.ops.trtllm, + "cute_dsl_bf16_bmm_locality_domain_inplace_rubin", + _record("locality_domain"), + ) + monkeypatch.setattr( + torch.ops.trtllm, + "cute_dsl_bf16_bmm_rubin", + _record("rubin"), + ) + monkeypatch.setattr( + torch.ops.trtllm, + "cute_dsl_bf16_bmm_blackwell", + _record("blackwell"), + ) + monkeypatch.setattr(torch.ops.trtllm, "bmm_out", _record("bmm_out")) + + a = torch.empty((2, 3, 4), dtype=torch.bfloat16) + weight = torch.empty((2, 6, 4), dtype=torch.bfloat16) + output = torch.empty((2, 3, 6), dtype=torch.bfloat16) + shards = (weight[:, :3].clone(), weight[:, 3:].clone()) if with_shards else None + + mla._bmm_bf16_out(a, weight, weight.transpose(1, 2), output, shards) + + assert len(calls) == 1 + op_name, op_args = calls[0] + assert op_name == expected_op + assert op_args[0] is a + assert op_args[-1] is output + if expected_op == "locality_domain": + assert shards is not None + assert op_args[1] is shards[0] + assert op_args[2] is shards[1] + elif expected_op in ("rubin", "blackwell"): + assert op_args[1] is weight diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index f65816f32a6f..551edab14244 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -13,6 +13,7 @@ from torch import nn from tensorrt_llm._torch import memory as memory_mod +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( @@ -620,3 +621,23 @@ def test_gms_unexpected_lock_state_raises(monkeypatch): backend.finalize_write.assert_not_called() backend.materialize_module.assert_not_called() checkpoint_loader.load_weights.assert_not_called() + + +def test_gms_rejects_localized_weights_before_model_init(monkeypatch): + events = [] + loader = _make_loader(monkeypatch, events=events) + loader._load_and_validate_config.return_value = SimpleNamespace( + mapping=SimpleNamespace(), + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ) + checkpoint_loader = MagicMock(name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "HF" + + with pytest.raises( + ValueError, + match="LoadFormat.GMS is incompatible with locality domain localized weights", + ): + loader.load("/ckpt", checkpoint_loader) + + model_loader_mod.AutoModelForCausalLM.from_config.assert_not_called() + checkpoint_loader.load_weights.assert_not_called() diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index b9f3fbacc2c8..630f4598f814 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -23,6 +23,7 @@ from tensorrt_llm._torch import distributed as distributed_mod from tensorrt_llm._torch.attention import mla as mla_mod from tensorrt_llm._torch.attention.mla import MLA +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models import modeling_llama as modeling_llama_mod from tensorrt_llm._torch.models import modeling_mistral as modeling_mistral_mod @@ -1781,11 +1782,29 @@ def test_speculative_mode_name_is_canonical_and_fails_closed( @pytest.mark.cpu_only +@pytest.mark.parametrize( + "qualification_case", + ["unregistered-model", "locality_domain-feature"], +) def test_mx_post_transform_receiver_falls_back_for_unqualified_model( - monkeypatch: pytest.MonkeyPatch, -) -> None: + monkeypatch, + qualification_case, +): events = [] loader = _make_loader(monkeypatch, events=events) + if qualification_case == "locality_domain-feature": + monkeypatch.setattr( + ModelLoader, + "_POST_TRANSFORM_PROFILE_REGISTRY", + _tiny_profile_registry(), + ) + model = _TinyModel(events) + model.model_config.locality_domain_policy = LocalityDomainPolicy(enabled=True) + monkeypatch.setattr( + model_loader_mod.AutoModelForCausalLM, + "from_config", + MagicMock(return_value=model), + ) checkpoint_loader = _PostTransformMxLoader(post_transform=True) model, _ = loader.load("/ckpt", checkpoint_loader) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py index 201d22142c6a..52f9ea8e7571 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv3.py @@ -21,12 +21,19 @@ """ from types import SimpleNamespace +from unittest.mock import patch import pytest +import torch +from tensorrt_llm._torch.locality_domain import policy +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models import modeling_deepseekv3, modeling_utils from tensorrt_llm._torch.models.modeling_deepseekv3 import Deepseekv3MoE +from tensorrt_llm._torch.modules import linear from tensorrt_llm._torch.moe.fused_moe import MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE +from tensorrt_llm._torch.peft.lora.layer import LoraLayer, LoraModuleType from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -35,6 +42,209 @@ EXPERTS_KEY = "model.layers.{}.mlp.experts" +@pytest.mark.parametrize("enable_locality_domains", [False, True]) +def test_linear_reassigned_quant_config_guard(enable_locality_domains: bool) -> None: + module = linear.Linear( + 16, + 16, + bias=False, + dtype=torch.bfloat16, + quant_config=QuantConfig(), + locality_domain_policy=policy.LocalityDomainPolicy(enabled=enable_locality_domains), + ) + weight = module.weight + module.quant_config = QuantConfig() + if enable_locality_domains: + with pytest.raises(RuntimeError, match="quant_config changed"): + module.create_weights() + else: + module.create_weights() + assert module.weight is weight + + +def test_layerwise_quant_config_preserves_unquantized_router() -> None: + gate = modeling_deepseekv3.DeepseekV3Gate( + hidden_size=16, + num_experts=4, + top_k=2, + n_group=2, + topk_group=1, + routed_scaling_factor=1.0, + dtype=torch.bfloat16, + ) + projection = linear.Linear(16, 16, bias=False, skip_create_weights_in_init=True) + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + model = torch.nn.Module() + model.add_module("gate", gate) + model.add_module("projection", projection) + model.model_config = SimpleNamespace( + quant_config_dict={"gate": quant_config, "projection": quant_config} + ) + + modeling_utils.apply_layerwise_quant_config(model.model_config, model.named_modules()) + gate.create_weights() + + assert gate.quant_config is None + assert gate.weight.dtype == torch.bfloat16 + assert projection.quant_config is quant_config + + +@pytest.mark.parametrize("use_cute_dsl_bf16_gemm", [False, True]) +@pytest.mark.parametrize("attach_lora", [False, True]) +@pytest.mark.parametrize("localized_weights", [False, True]) +def test_deepseek_small_m_selects_kernel_for_localized_weights( + use_cute_dsl_bf16_gemm: bool, + attach_lora: bool, + localized_weights: bool, +) -> None: + module = modeling_deepseekv3.DeepseekV3Linear( + 16, + 16, + bias=False, + dtype=torch.bfloat16, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + lora=LoraLayer([LoraModuleType.ATTENTION_Q], [16]) if attach_lora else None, + ) + if localized_weights: + module._locality_domain_weight_shards = tuple(module.weight.detach().chunk(2)) + module.weight = torch.nn.Parameter( + torch.empty(0, dtype=torch.bfloat16), requires_grad=False + ) + inputs = torch.ones(1, 16, dtype=torch.bfloat16) + expected = torch.ones_like(inputs) + with ( + patch.object(modeling_deepseekv3, "get_sm_version", return_value=107), + patch.object(modeling_deepseekv3, "is_sm_100f", return_value=True), + patch.object(torch.ops.trtllm, "dsv3_fused_a_gemm_op", return_value=expected) as fused, + patch.object(linear.Linear, "apply_linear", return_value=expected) as fallback, + ): + assert module.apply_linear(inputs, None, layer_idx=3) is expected + if localized_weights: + fallback.assert_called_once_with(inputs, None, None, 3) + fused.assert_not_called() + else: + fused.assert_called_once() + fallback.assert_not_called() + + +@pytest.mark.parametrize("num_tokens", [1, 16]) +@pytest.mark.parametrize("use_cute_dsl_bf16_gemm", [False, True]) +def test_deepseek_small_m_preserves_lora(num_tokens: int, use_cute_dsl_bf16_gemm: bool) -> None: + lora_layer = LoraLayer([LoraModuleType.ATTENTION_Q], [16]) + module = modeling_deepseekv3.DeepseekV3Linear( + 16, + 16, + dtype=torch.bfloat16, + lora=lora_layer, + use_cute_dsl_bf16_gemm=use_cute_dsl_bf16_gemm, + ) + module.weight.data.copy_(2 * torch.eye(16, dtype=torch.bfloat16)) + module.bias.data.fill_(1) + inputs = torch.ones(num_tokens, 16, dtype=torch.bfloat16) + adapter_output = torch.full_like(inputs, 3) + ranks = torch.tensor([1], dtype=torch.int32) + weight_pointers = torch.zeros((1, 2), dtype=torch.int64) + lora_params = { + "num_seqs": 1, + "host_request_types": torch.zeros(1, dtype=torch.int32), + "prompt_lens_cpu": torch.tensor([num_tokens], dtype=torch.int32), + 3: { + int(LoraModuleType.ATTENTION_Q): { + "adapter_size": ranks, + "weight_pointers": weight_pointers, + } + }, + } + + # Keep Linear.apply_linear and the LoRA merge real. Use CPU implementations + # for the fused base GEMM and the native adapter GEMM. + with ( + patch.object(modeling_deepseekv3, "get_sm_version", return_value=107), + patch.object(linear, "is_sm_100f", return_value=False), + patch.object( + torch.ops.trtllm, + "dsv3_fused_a_gemm_op", + side_effect=lambda x, weight, bias, _: x @ weight + bias, + ) as fused, + patch.object( + torch.ops.trtllm, "lora_grouped_gemm", return_value=adapter_output + ) as adapter_gemm, + ): + output = module.apply_linear(inputs, module.bias, lora_params, layer_idx=3) + + torch.testing.assert_close(output, torch.full_like(inputs, 6)) + fused.assert_not_called() + adapter_gemm.assert_called_once() + assert adapter_gemm.call_args.args[2][0] is ranks + assert adapter_gemm.call_args.args[3][0] is weight_pointers + + +@pytest.mark.parametrize("quant_algo", [None, QuantAlgo.NVFP4], ids=["bf16", "nvfp4"]) +@pytest.mark.parametrize("per_module_quant", [False, True]) +@pytest.mark.parametrize("frozen_config", [False, True]) +def test_shared_experts_do_not_partition_locality_domains( + quant_algo, per_module_quant: bool, frozen_config: bool +) -> None: + shared_quant_config = QuantConfig(quant_algo=quant_algo, group_size=16) + global_quant_config = ( + QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION) + if per_module_quant + else shared_quant_config + ) + model_policy = policy.LocalityDomainPolicy(enabled=True) + model_config = ModelConfig( + pretrained_config=SimpleNamespace(n_group=1, topk_group=1, routed_scaling_factor=1.0), + quant_config=global_quant_config, + quant_config_dict=( + {"model.layers.0.mlp.shared_experts": shared_quant_config} if per_module_quant else None + ), + skip_create_weights_in_init=True, + use_cute_dsl_bf16_gemm=True, + locality_domain_policy=model_policy, + ) + model_config._frozen = frozen_config + + # Exercise the real GatedMLP/Linear construction and planner on CPU, while + # presenting hardware on which both BF16 and NVFP4 can be partitioned. + with ( + patch( + "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", + return_value=True, + ), + patch("tensorrt_llm._torch.cute_dsl_utils.IS_CUTLASS_DSL_AVAILABLE", True), + patch("tensorrt_llm._torch.cute_dsl_utils.IS_CUTLASS_DSL_RUBIN_AVAILABLE", True), + patch.object( + modeling_deepseekv3, "create_moe", return_value=torch.nn.Module() + ) as create_moe, + patch.object(torch.cuda, "is_available", return_value=False), + patch.object(torch.cuda, "Event"), + ): + module = Deepseekv3MoE( + num_experts=4, + top_k=2, + hidden_size=256, + intermediate_size=256, + shared_expert_intermediate_size=256, + aux_stream_dict={modeling_deepseekv3.AuxStreamType.MoeShared: None}, + dtype=torch.bfloat16, + model_config=model_config, + layer_idx=0, + ) + for projection in (module.shared_experts.gate_up_proj, module.shared_experts.down_proj): + projection.create_weights() + assert projection.quant_config is shared_quant_config + assert not projection.partition_plan.enabled + assert projection._locality_domain_runtime is None + + assert module.shared_experts_use_fp4 == (quant_algo == QuantAlgo.NVFP4) + assert create_moe.call_args.kwargs["model_config"] is model_config + assert model_config.quant_config is global_quant_config + assert model_config.locality_domain_policy is model_policy + assert model_config.extra_attrs["locality_domain_policy"] is model_policy + assert model_config.locality_domain_policy.enabled + assert model_config._frozen == frozen_config + + @pytest.fixture def mixed_precision_config(): return QuantConfig(quant_algo=QuantAlgo.MIXED_PRECISION) diff --git a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py index f940d5fe7910..a6f43f942671 100644 --- a/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py +++ b/tests/unittest/_torch/modeling/test_modeling_deepseekv4.py @@ -275,6 +275,59 @@ def test_deepseek_v4_eplb_weight_loader_pages_out_each_moe_layer(monkeypatch): assert pageout_calls == [(".safetensors", "dontneed")] +@pytest.mark.cpu_only +@pytest.mark.parametrize("partition_enabled", [False, True]) +def test_deepseek_v4_loader_rebinds_absorption_weight(partition_enabled: bool) -> None: + from tensorrt_llm._torch.attention import mla as mla_module + from tensorrt_llm.models import modeling_utils + + attention = mla_module.MLA.__new__(mla_module.MLA) + nn.Module.__init__(attention) + attention.kv_b_proj = nn.Linear(4, 16, bias=False, dtype=torch.bfloat16) + attention.kv_b_proj.partition_plan = SimpleNamespace( + enabled=partition_enabled, op_kind="bf16_linear" + ) + attention.k_b_proj_trans = nn.Parameter(torch.zeros(2, 4, 4, dtype=torch.bfloat16)) + attention._weights_transformed = True + + model = nn.Module() + model.model = nn.Module() + model.model.layers = nn.ModuleList([nn.Module()]) + model.model.layers[0].self_attn = attention + model.config = SimpleNamespace( + q_lora_rank=4, + num_attention_heads=2, + qk_nope_head_dim=4, + v_head_dim=4, + kv_lora_rank=4, + num_hidden_layers=1, + num_nextn_predict_layers=0, + ) + model.model_config = SimpleNamespace( + mapping=SimpleNamespace( + tp_rank=0, + tp_size=1, + cp_rank=0, + cp_size=1, + enable_attention_dp=False, + ), + moe_load_balancer=None, + quant_config=modeling_utils.QuantConfig(), + ) + weight = torch.arange(64, dtype=torch.bfloat16).view(16, 4) + DeepseekV4WeightLoader(model)._load_weights_impl( + {"model.layers.0.self_attn.kv_b_proj.weight": weight} + ) + + torch.testing.assert_close(attention.v_b_proj, weight.view(2, 8, 4)[:, 4:]) + shares_storage = ( + attention.v_b_proj.untyped_storage().data_ptr() + == attention.kv_b_proj.weight.untyped_storage().data_ptr() + ) + assert shares_storage is not partition_enabled + assert not attention._weights_transformed + + def test_deepseek_v4_fused_a_weight_scale_rebuilds_fp8_shape(): module = torch.nn.Module() module.weight = torch.nn.Parameter(torch.empty((2048, 7168), dtype=torch.float8_e4m3fn)) diff --git a/tests/unittest/_torch/modules/test_gated_mlp.py b/tests/unittest/_torch/modules/test_gated_mlp.py index da0b31ee3b1f..5cbf94924c82 100644 --- a/tests/unittest/_torch/modules/test_gated_mlp.py +++ b/tests/unittest/_torch/modules/test_gated_mlp.py @@ -10,8 +10,12 @@ import torch.nn.functional as F from torch import nn +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy +from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules import gated_mlp as gated_mlp_module from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import Linear as _RealLinear +from tensorrt_llm._torch.visual_gen import config as visual_gen_config def _make_gate_up_proj( @@ -40,6 +44,25 @@ def _make_down_proj() -> nn.Module: return down_proj +@pytest.mark.cpu_only +def test_diffusion_config_without_locality_domain_policy() -> None: + config = visual_gen_config.DiffusionModelConfig() + mlp = GatedMLP( + hidden_size=8, + intermediate_size=16, + bias=False, + dtype=torch.bfloat16, + config=config, + ) + + assert mlp.gate_up_proj.weight.shape == (32, 8) + assert mlp.down_proj.weight.shape == (8, 16) + for projection in (mlp.gate_up_proj, mlp.down_proj): + assert not projection._locality_domain_policy.enabled + assert not projection.partition_plan.enabled + assert projection._locality_domain_weight_shards is None + + def test_gate_up_partition_falls_back_to_swiglu( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -120,3 +143,103 @@ def test_activation_controls_fp8_quant_fusion_capability( monkeypatch.setattr(gated_mlp_module, "IS_CUTLASS_DSL_RUBIN_AVAILABLE", True) assert mlp._can_fuse_swiglu_fp8_quant() is expected + + +@pytest.mark.parametrize( + ("mlp_kwargs", "expected_enabled"), + [ + pytest.param({}, False, id="default"), + pytest.param( + {"enable_locality_domain_bf16_linear": True}, + True, + id="explicit-opt-in", + ), + ], +) +def test_locality_domain_bf16_linear_opt_in_propagates_to_both_projections( + monkeypatch: pytest.MonkeyPatch, + mlp_kwargs: dict[str, bool], + expected_enabled: bool, +) -> None: + policy = LocalityDomainPolicy(enabled=True) + config = ModelConfig( + use_cute_dsl_bf16_gemm=True, + locality_domain_policy=policy, + ) + linear_constructor = Mock(side_effect=[nn.Module(), nn.Module()]) + # GatedMLP.__init__ calls Linear._calc_shard for uneven-TP sharding; + # the Mock must delegate to the real staticmethod or the arithmetic + # downstream operates on Mock objects. + linear_constructor._calc_shard = _RealLinear._calc_shard + monkeypatch.setattr(gated_mlp_module, "Linear", linear_constructor) + + GatedMLP( + hidden_size=8, + intermediate_size=16, + bias=False, + dtype=torch.bfloat16, + config=config, + **mlp_kwargs, + ) + + assert linear_constructor.call_count == 2 + gate_up_call, down_call = linear_constructor.call_args_list + assert gate_up_call.args == (8, 32) + assert down_call.args == (16, 8) + for projection_call in (gate_up_call, down_call): + assert projection_call.kwargs["use_cute_dsl_bf16_gemm"] is True + assert projection_call.kwargs["enable_locality_domain_bf16_linear"] is expected_enabled + assert projection_call.kwargs["locality_domain_policy"] is policy + + +def test_deepseek_dense_layer_enables_locality_domain_bf16_mlp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.models import modeling_deepseekv3 + from tensorrt_llm._torch.utils import AuxStreamType + from tensorrt_llm.models.modeling_utils import QuantConfig + + policy = LocalityDomainPolicy(enabled=True) + pretrained_config = SimpleNamespace( + hidden_size=128, + moe_intermediate_size=64, + n_routed_experts=None, + n_shared_experts=0, + num_experts_per_tok=1, + model_type="deepseek_v3", + first_k_dense_replace=3, + moe_layer_freq=1, + intermediate_size=256, + torch_dtype=torch.bfloat16, + rms_norm_eps=1e-6, + ) + config = ModelConfig( + pretrained_config=pretrained_config, + quant_config=QuantConfig(), + skip_create_weights_in_init=True, + use_cute_dsl_bf16_gemm=True, + locality_domain_policy=policy, + ) + monkeypatch.setattr( + modeling_deepseekv3, + "DeepseekV3Attention", + Mock(return_value=nn.Identity()), + ) + monkeypatch.setattr( + modeling_deepseekv3, + "can_access_peer", + Mock(return_value=False), + ) + gated_mlp_constructor = Mock(return_value=nn.Identity()) + monkeypatch.setattr(modeling_deepseekv3, "GatedMLP", gated_mlp_constructor) + + modeling_deepseekv3.DeepseekV3DecoderLayer( + config, + layer_idx=0, + aux_stream_dict={AuxStreamType.Attention: None}, + ) + + assert gated_mlp_constructor.call_count == 1 + call = gated_mlp_constructor.call_args + assert call.kwargs["config"] is config + assert call.kwargs["enable_locality_domain_bf16_linear"] is True diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index e195d6f9d973..0e55ea512029 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -178,7 +178,13 @@ def test_fp8_block_scale_moe_fallback_tactic_is_explicit_and_deterministic(): _select_explicit_fallback_tactic([]) -def test_cutedsl_count_native_runner_keeps_both_tactics_and_threads_counts() -> None: +@pytest.mark.parametrize("use_direct_expert_metadata", [False, True]) +def test_cutedsl_count_native_runner_keeps_both_tactics_and_threads_counts( + monkeypatch: pytest.MonkeyPatch, use_direct_expert_metadata: bool +) -> None: + monkeypatch.setattr( + "tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl.get_sm_version", lambda: 100 + ) forward_impl = MagicMock(return_value=torch.empty(1)) runner = CuteDslFusedMoENvfp4Runner( forward_impl=forward_impl, @@ -186,19 +192,27 @@ def test_cutedsl_count_native_runner_keeps_both_tactics_and_threads_counts() -> top_k=1, num_local_experts=8, local_expert_offset=0, - use_direct_expert_metadata=True, + use_direct_expert_metadata=use_direct_expert_metadata, ) counts = torch.tensor([0, 1, 31, 32, 7, 0, 16, 2], dtype=torch.int32) - inputs = [torch.empty(8 * 32, 1) for _ in range(6)] + [counts] + inputs = [torch.empty(8 * 32, 1) for _ in range(6)] + forward_inputs = inputs + [counts] if use_direct_expert_metadata else inputs - runner.forward(inputs, tactic=256) + assert runner.forward(forward_inputs, tactic=256) is forward_impl.return_value - assert runner.unique_id()[-1] == "direct_expert_metadata" assert runner.get_valid_tactics([], OptimizationProfile()) == [128, 256] - assert runner.get_tuning_config().inputs_pre_hook is None - assert forward_impl.call_args.kwargs["recv_expert_count"] is counts - assert forward_impl.call_args.kwargs["deep_ep_expert_capacity"] == 32 - assert forward_impl.call_args.kwargs["use_count_native_expert_metadata"] is True + expected_kwargs = {"enable_alltoall": False, "tile_size": 256} + if use_direct_expert_metadata: + assert runner.unique_id()[-1] == "direct_expert_metadata" + assert runner.get_tuning_config().inputs_pre_hook is None + expected_kwargs.update( + recv_expert_count=counts, + deep_ep_expert_capacity=32, + use_count_native_expert_metadata=True, + ) + else: + assert runner.get_tuning_config().inputs_pre_hook is not None + forward_impl.assert_called_once_with(*inputs, **expected_kwargs) def test_cutedsl_direct_metadata_tuning_buckets_capacity_and_legacy_uses_tokens() -> None: diff --git a/tests/unittest/_torch/multi_gpu/test_linear.py b/tests/unittest/_torch/multi_gpu/test_linear.py index 4db1c6b053ce..9da57114d9cb 100644 --- a/tests/unittest/_torch/multi_gpu/test_linear.py +++ b/tests/unittest/_torch/multi_gpu/test_linear.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import pickle import sys import traceback @@ -12,8 +15,13 @@ import tensorrt_llm import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.autotuner import autotune +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy +from tensorrt_llm._torch.locality_domain_utils import is_locality_domain_enabled from tensorrt_llm._torch.modules.linear import Linear, TensorParallelMode -from tensorrt_llm.functional import AllReduceFusionOp, AllReduceParams +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.functional import (AllReduceFusionOp, AllReduceParams, + AllReduceStrategy) from tensorrt_llm.mapping import Mapping from tensorrt_llm.math_utils import pad_up from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -344,10 +352,20 @@ def check_accuracy(a, b, atol, rtol, percent): @torch.inference_mode -def fp4_row_linear_allreduce(tp_size, local_rank, seq_len, output_size, - hidden_size, dtype, output_ref, x_sf_global, - w_sf_global, x_fp4s, w_fp4, x_sf_blocks, - w_sf_block_unswizzled): +def fp4_row_linear_allreduce(tp_size, + local_rank, + seq_len, + output_size, + hidden_size, + dtype, + output_ref, + x_sf_global, + w_sf_global, + x_fp4s, + w_fp4, + x_sf_blocks, + w_sf_block_unswizzled, + use_locality_domain=False): output_ref = output_ref.cuda() x_sf_global = x_sf_global.cuda() w_sf_global = w_sf_global.cuda() @@ -369,6 +387,9 @@ def fp4_row_linear_allreduce(tp_size, local_rank, seq_len, output_size, rank=local_rank, ), tensor_parallel_mode=TensorParallelMode.ROW, + allreduce_strategy=AllReduceStrategy.AUTO, + locality_domain_policy=LocalityDomainPolicy( + enabled=use_locality_domain), ) l0.load_weights([{ @@ -383,28 +404,54 @@ def fp4_row_linear_allreduce(tp_size, local_rank, seq_len, output_size, }]) l0.cuda() + if use_locality_domain: + l0.post_load_weights() + assert l0.partition_plan.enabled + assert l0._locality_domain_weight_shards is not None + assert len(l0._locality_domain_weight_shards) == 2 + assert l0.weight.numel() == 0 + assert l0.weight_scale.numel() == 0 + assert l0.all_reduce.strategy == AllReduceStrategy.AUTO + assert l0.all_reduce.uses_nccl_symmetric_memory_window() + # Guarantee that this regression exercises the routing guard even when + # fused GEMM-allreduce is unavailable in a particular test environment. + l0.use_fused_gemm_allreduce = True + # TODO: parameters['weight']' size mismatch at index 0 # l0 = torch.compile(l0) with torch.inference_mode(), autotune(): - output = l0.forward((x_fp4, x_sf_block)) + if use_locality_domain: + output = l0.forward((x_fp4, x_sf_block), + all_reduce_params=AllReduceParams()) + else: + output = l0.forward((x_fp4, x_sf_block)) torch.cuda.synchronize() - check_accuracy(output, output_ref, atol=0.05, rtol=0.05, percent=0.99) - - -def fp4_row_linear_allreduce_run_single_rank(func, tp_size, seq_len, - output_size, hidden_size, dtype, - output_ref, x_sf_global, - w_sf_global, x_fp4s, w_fp4, + atol = 0.15 if use_locality_domain else 0.05 + check_accuracy(output, output_ref, atol=atol, rtol=0.05, percent=0.99) + + +def fp4_row_linear_allreduce_run_single_rank(func, + tp_size, + seq_len, + output_size, + hidden_size, + dtype, + output_ref, + x_sf_global, + w_sf_global, + x_fp4s, + w_fp4, x_sf_blocks, - w_sf_block_unswizzled): + w_sf_block_unswizzled, + use_locality_domain=False): local_rank = tensorrt_llm.mpi_rank() torch.cuda.set_device(local_rank) try: func(tp_size, local_rank, seq_len, output_size, hidden_size, dtype, output_ref, x_sf_global, w_sf_global, x_fp4s, w_fp4, x_sf_blocks, - w_sf_block_unswizzled) + w_sf_block_unswizzled, use_locality_domain) except Exception as e: print(f"Error: {e}") raise @@ -475,3 +522,199 @@ def test_fp4_row_linear_allreduce(seq_len, output_size, hidden_size, dtype, for r in results: assert r is True + + +@torch.inference_mode +def bf16_row_linear_locality_domain_forward(x, hidden_size, dtype, + tensor_parallel_size, + tensor_parallel_rank, weights): + x = x.cuda() + output_size = weights[0].shape[0] + local_hidden_size = hidden_size // tensor_parallel_size + + linear = Linear( + in_features=hidden_size, + out_features=output_size, + bias=False, + dtype=dtype, + mapping=Mapping( + world_size=tensor_parallel_size, + tp_size=tensor_parallel_size, + rank=tensor_parallel_rank, + ), + tensor_parallel_mode=TensorParallelMode.ROW, + allreduce_strategy=AllReduceStrategy.AUTO, + use_cute_dsl_bf16_gemm=True, + enable_locality_domain_bf16_linear=True, + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ) + linear.load_weights([{"weight": weights[0]}]) + linear.cuda() + linear.post_load_weights() + + assert linear.partition_plan.enabled + assert linear.partition_plan.op_kind == "bf16_linear" + assert linear._locality_domain_weight_shards is not None + assert len(linear._locality_domain_weight_shards) == 2 + assert linear.weight.numel() == 0 + assert linear._locality_domain_weight_shards[0]["weight"].shape == ( + output_size // 2, + local_hidden_size, + ) + # ROW-TP must not select a path that reads the released full weight even if + # a fused GEMM/all-reduce capability flag is set. + linear.use_fused_gemm_allreduce = True + + input_shard = torch.chunk(x, tensor_parallel_size, + dim=-1)[tensor_parallel_rank].contiguous() + with autotune(): + output = linear(input_shard, all_reduce_params=AllReduceParams()) + expected = x.float() @ weights[0].cuda().t().float() + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=0.5) + + +def _prepare_fp4_row_linear_allreduce_case( + seq_len: int, + output_size: int, + hidden_size: int, + dtype: torch.dtype, + tp_size: int, + sample_device: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor], + torch.Tensor, list[torch.Tensor], torch.Tensor]: + scaling_vector_size = 16 + x = torch.randn((seq_len, hidden_size), dtype=dtype, + device=sample_device).cuda() + weight = torch.randn((output_size, hidden_size), + dtype=dtype, + device=sample_device).cuda() + + x_sf_global = (448 * 6) / x.abs().max().float() + x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize(x, x_sf_global, + scaling_vector_size, + False) + weight_sf_global = (448 * 6) / weight.abs().max().float() + weight_fp4, weight_sf_block = torch.ops.trtllm.fp4_quantize( + weight, weight_sf_global, scaling_vector_size, False) + weight_sf_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + weight_sf_block.cpu().view(pad_up(output_size, 128), -1)) + + with torch.inference_mode(): + alpha_ref = 1.0 / (weight_sf_global * x_sf_global) + output_ref = torch.ops.trtllm.fp4_gemm( + x_fp4, + weight_fp4, + x_sf_block, + weight_sf_block, + alpha_ref, + fp4_utils.FP4GemmType.W4A4_NVFP4_NVFP4, + dtype, + ) + + input_fp4_shards = [] + input_sf_shards = [] + for input_shard in torch.chunk(x, tp_size, dim=-1): + input_fp4, input_sf = torch.ops.trtllm.fp4_quantize( + input_shard.contiguous(), x_sf_global, scaling_vector_size, False) + input_fp4_shards.append(input_fp4.cpu()) + input_sf_shards.append(input_sf.cpu()) + + torch.cuda.synchronize() + return ( + output_ref.cpu(), + x_sf_global.cpu(), + weight_sf_global.cpu(), + input_fp4_shards, + weight_fp4.cpu(), + input_sf_shards, + weight_sf_unswizzled.cpu(), + ) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="CuTe DSL package with Rubin support is not available", +) +@pytest.mark.skipif(torch.cuda.device_count() < 2, + reason="needs 2 GPUs to run this test") +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_fp4_row_linear_allreduce_locality_domain_freed_weights( + mpi_pool_executor, monkeypatch): + is_locality_domain_enabled.cache_clear() + if not is_locality_domain_enabled(): + pytest.skip( + "locality domain localization is not enabled/supported on this system" + ) + + monkeypatch.setenv("TRTLLM_GEMM_ALLREDUCE_FUSION_ENABLED", "1") + monkeypatch.setenv("TLLM_NCCL_SYMMETRIC_ZERO_COPY", "1") + torch.manual_seed(42) + + seq_len = 1 + output_size = 1536 + hidden_size = 4224 + dtype = torch.bfloat16 + tp_size = mpi_pool_executor.num_workers + case_inputs = _prepare_fp4_row_linear_allreduce_case(seq_len, + output_size, + hidden_size, + dtype, + tp_size, + sample_device="cuda") + rank_args = (fp4_row_linear_allreduce, tp_size, seq_len, output_size, + hidden_size, dtype, *case_inputs, True) + results = mpi_pool_executor.map( + fp4_row_linear_allreduce_run_single_rank, + *zip(*[rank_args] * tp_size), + ) + + for result in results: + assert result is True + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="CuTe DSL package with Rubin support is not available", +) +@pytest.mark.skipif(torch.cuda.device_count() < 2, + reason="needs 2 GPUs to run this test") +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_bf16_row_linear_allreduce_locality_domain_freed_weights( + mpi_pool_executor): + is_locality_domain_enabled.cache_clear() + if not is_locality_domain_enabled(): + pytest.skip( + "locality domain localization is not enabled/supported on this system" + ) + + tensor_parallel_size = mpi_pool_executor.num_workers + torch.manual_seed(107) + seq_len = 1 + hidden_size = 512 + output_size = 256 + dtype = torch.bfloat16 + x = torch.randn((seq_len, hidden_size), dtype=dtype) + weight = torch.randn((output_size, hidden_size), dtype=dtype) + + results = mpi_pool_executor.map( + run_single_rank, + *zip(*[( + tensor_parallel_size, + bf16_row_linear_locality_domain_forward, + x, + [weight], + hidden_size, + dtype, + )] * tensor_parallel_size), + ) + for result in results: + assert result is True diff --git a/tests/unittest/_torch/thop/parallel/_cute_dsl_bf16_rubin_test_utils.py b/tests/unittest/_torch/thop/parallel/_cute_dsl_bf16_rubin_test_utils.py new file mode 100644 index 000000000000..16b69309d7ab --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/_cute_dsl_bf16_rubin_test_utils.py @@ -0,0 +1,194 @@ +# 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. + +import math +from collections.abc import Sequence +from typing import Literal + +import pytest +import torch + +from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile, TunableRunner +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, +) +from tensorrt_llm._torch.locality_domain.autotune import LocalityDomainConcurrentTunableRunner +from tensorrt_llm._torch.locality_domain_utils import is_locality_domain_enabled +from tensorrt_llm._utils import get_sm_version + +KernelVariant = Literal["base", "preferred_cluster"] +Tactic = tuple[object, ...] + +RUBIN_CUTE_DSL_MARKS = [ + pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", + ), + pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available"), + pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="CuTe DSL package with Rubin support is not available", + ), +] + + +def skip_if_no_locality_domain() -> None: + is_locality_domain_enabled.cache_clear() + if not is_locality_domain_enabled(): + pytest.skip("locality domain localization is not enabled/supported on this system") + + +def reset_bf16_gemm_state() -> None: + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + AutoTuner.get().clear_cache() + runner_class = cute_dsl_custom_ops.CuteDSLBf16RubinGemmRunner + runner_class.kernel_cache.clear() + runner_class.split_k_gemm_cache.clear() + + +def make_bf16_gemm_runner() -> TunableRunner: + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + reset_bf16_gemm_state() + return cute_dsl_custom_ops.CuteDSLBf16RubinGemmRunner(use_tvm_ffi=True) + + +def make_bf16_bmm_runner() -> TunableRunner: + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + AutoTuner.get().clear_cache() + runner_class = cute_dsl_custom_ops.CuteDSLBf16RubinBmmRunner + runner_class.kernel_cache.clear() + return runner_class(use_tvm_ffi=True) + + +def select_bf16_tactic( + tactics: Sequence[Tactic], + kernel_variant: KernelVariant, + *, + split_k_slices: int | None = None, +) -> Tactic: + candidates = [ + tactic for tactic in tactics if tactic[0] == kernel_variant and tactic[1] is False + ] + if split_k_slices is not None and kernel_variant == "base": + candidates = [ + tactic + for tactic in candidates + if len(tactic) == 6 and tactic[0] == "base" and tactic[-1] == split_k_slices + ] + elif split_k_slices not in (None, 1): + candidates = [] + + tactic_description = kernel_variant + if split_k_slices is not None: + tactic_description += f" with split_k_slices={split_k_slices}" + assert candidates, f"no {tactic_description} tactic found" + return candidates[0] + + +def select_captured_locality_domain_tactic( + tactics_capture: AutoTuner.TacticsCapture, + op_name: str, + kernel_variant: KernelVariant, + *, + split_k_slices: int | None = None, +) -> tuple[LocalityDomainConcurrentTunableRunner, Tactic, list[Tactic]]: + assert len(tactics_capture._captured_contexts) == 1 + context = tactics_capture._captured_contexts[0] + assert context["custom_op"] == f"trtllm::{op_name}::locality_domain_concurrent" + assert len(context["runners"]) == 1 + concurrent_runner = context["runners"][0] + assert isinstance(concurrent_runner, LocalityDomainConcurrentTunableRunner) + tactics = concurrent_runner.get_valid_tactics(context["inputs"], OptimizationProfile()) + tactic = select_bf16_tactic( + tactics, + kernel_variant, + split_k_slices=split_k_slices, + ) + return concurrent_runner, tactic, tactics + + +def bf16_matmul_atol(k: int) -> float: + """Return the absolute tolerance for BF16 matmuls with unit-variance inputs. + + Args: + k: Reduction dimension of the GEMM or BMM. + """ + return 4e-3 * math.sqrt(k) + + +def run_locality_domain_composite( + op_name: str, + args: tuple[torch.Tensor, ...], + expected_partitions: Sequence[torch.Tensor], + partition_dim: int, + kernel_variant: KernelVariant, + *, + split_k_slices: int | None = None, + capture_graph: bool = True, + atol: float | None = None, +) -> list[Tactic]: + op = getattr(torch.ops.trtllm, op_name) + output = args[-1] + tuner = AutoTuner.get() + tuner.clear_cache() + with tuner.capture() as tactics_capture: + op(*args) + + concurrent_runner, tactic, tactics = select_captured_locality_domain_tactic( + tactics_capture, + op_name, + kernel_variant, + split_k_slices=split_k_slices, + ) + context = tactics_capture._captured_contexts[0] + + if capture_graph: + graph = torch.cuda.CUDAGraph() + with tuner.replay(((concurrent_runner, tactic),)): + # Compile both partition-local artifacts before capture. + op(*args) + torch.cuda.synchronize() + with torch.cuda.graph(graph): + op(*args) + + output.fill_(float("nan")) + graph.replay() + else: + output.fill_(float("nan")) + concurrent_runner(context["inputs"], tactic=tactic) + torch.cuda.synchronize() + + # Scale the absolute allowance with K for these unit-variance inputs. + if atol is None: + atol = bf16_matmul_atol(args[0].shape[-1]) + output_partitions = output.chunk(len(expected_partitions), dim=partition_dim) + assert len(output_partitions) == len(expected_partitions) + for output_partition, expected in zip( + output_partitions, + expected_partitions, + strict=True, + ): + torch.testing.assert_close( + output_partition.float(), + expected.float(), + rtol=1e-2, + atol=atol, + ) + + return tactics diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py new file mode 100644 index 000000000000..ea8c337ed03e --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py @@ -0,0 +1,537 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from _torch.thop.parallel._cute_dsl_bf16_rubin_test_utils import ( + RUBIN_CUTE_DSL_MARKS, + bf16_matmul_atol, + make_bf16_bmm_runner, + make_bf16_gemm_runner, + reset_bf16_gemm_state, + run_locality_domain_composite, + select_bf16_tactic, + select_captured_locality_domain_tactic, + skip_if_no_locality_domain, +) + +from tensorrt_llm._torch.autotuner import AutoTuner, autotune + +pytestmark = RUBIN_CUTE_DSL_MARKS + + +@pytest.mark.parametrize( + ("kernel_variant", "mnk"), + [("base", (128, 128, 128)), ("preferred_cluster", (256, 256, 128))], +) +def test_cute_dsl_bf16_gemm_locality_domain_rubin(kernel_variant, mnk): + skip_if_no_locality_domain() + torch.manual_seed(123) + runner = make_bf16_gemm_runner() + + m, n, k = mnk + act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + weight_0 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + weight_1 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + output = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + tactics = runner.get_valid_tactics([act, weight_0, output], None) + tactic = select_bf16_tactic(tactics, kernel_variant, split_k_slices=1) + + expected_0 = act.float() @ weight_0.t().float() + expected_1 = act.float() @ weight_1.t().float() + runner([act, weight_0, output], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected_0, rtol=1e-2, atol=bf16_matmul_atol(k)) + + wide_output = torch.empty(m, n * 2, dtype=torch.bfloat16, device="cuda") + run_locality_domain_composite( + "cute_dsl_bf16_gemm_locality_domain_inplace_rubin", + (act, weight_0, weight_1, wide_output), + (expected_0, expected_1), + partition_dim=1, + kernel_variant=kernel_variant, + split_k_slices=1, + ) + + +def test_cute_dsl_bf16_gemm_locality_domain_production_shapes_rubin(): + """Exercise the exact per-partition shapes used by DeepSeek BF16 projections.""" + skip_if_no_locality_domain() + shapes = ( + ("fused_a", 1, 1056, 7168), + ("q_b", 1, 12288, 1536), + ("kv_b", 1, 16384, 512), + ("o_proj", 1, 3584, 16384), + ("gate_up", 1, 18432, 7168), + ("down", 1, 3584, 18432), + ) + + for workload, m, n, k in shapes: + torch.manual_seed(2112) + act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + weight_0 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + weight_1 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + output = torch.empty(m, n * 2, dtype=torch.bfloat16, device="cuda") + expected_0 = act.float() @ weight_0.t().float() + expected_1 = act.float() @ weight_1.t().float() + + try: + run_locality_domain_composite( + "cute_dsl_bf16_gemm_locality_domain_inplace_rubin", + (act, weight_0, weight_1, output), + (expected_0, expected_1), + partition_dim=1, + kernel_variant="base", + split_k_slices=1, + capture_graph=workload == "fused_a", + ) + except AssertionError as error: + raise AssertionError(f"{workload} shape {(m, n, k)} failed") from error + + +def test_deepseek_gate_bf16_locality_domain_end_to_end_rubin(): + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy + from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate + + skip_if_no_locality_domain() + torch.manual_seed(2114) + reset_bf16_gemm_state() + + hidden_states = torch.randn(1, 1, 7168, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(256, 7168, dtype=torch.bfloat16, device="cuda") + correction_bias = torch.randn(256, dtype=torch.float32, device="cuda") + expected = hidden_states.float() @ weight.t().float() + atol = bf16_matmul_atol(hidden_states.shape[-1]) + + gate = DeepseekV3Gate( + hidden_size=7168, + num_experts=256, + top_k=8, + n_group=8, + topk_group=4, + routed_scaling_factor=2.5, + dtype=torch.bfloat16, + use_cute_dsl_bf16_gemm=True, + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ).cuda() + gate.load_weights([{"weight": weight, "e_score_correction_bias": correction_bias}]) + gate.post_load_weights() + + assert gate.partition_plan.enabled + assert gate.partition_plan.op_kind == "bf16_linear" + assert gate._locality_domain_weight_shards is not None + assert gate.weight.numel() == 0 + torch.testing.assert_close(gate.e_score_correction_bias, correction_bias) + + tuner = AutoTuner.get() + with tuner.capture() as tactics_capture, torch.inference_mode(): + gate(hidden_states) + concurrent_runner, tactic, _ = select_captured_locality_domain_tactic( + tactics_capture, + "cute_dsl_bf16_gemm_locality_domain_inplace_rubin", + "base", + split_k_slices=1, + ) + + graph = torch.cuda.CUDAGraph() + with tuner.replay(((concurrent_runner, tactic),)), torch.inference_mode(): + gate(hidden_states) + torch.cuda.synchronize() + with torch.cuda.graph(graph): + output = gate(hidden_states) + with torch.inference_mode(): + output.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + + assert output.dtype == torch.float32 + assert output.shape == (1, 1, 256) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=atol) + + first_generation_shards = gate._locality_domain_weight_shards + gate.pre_reload_weights() + assert gate._locality_domain_weight_shards is None + assert gate.weight.shape == weight.shape + + reloaded_weight = torch.randn_like(weight) + reloaded_bias = torch.randn_like(correction_bias) + reloaded_expected = hidden_states.float() @ reloaded_weight.t().float() + gate.load_weights([{"weight": reloaded_weight, "e_score_correction_bias": reloaded_bias}]) + gate.post_load_weights() + + assert gate._locality_domain_weight_shards is not None + assert gate._locality_domain_weight_shards is not first_generation_shards + assert gate.weight.numel() == 0 + torch.testing.assert_close(gate.e_score_correction_bias, reloaded_bias) + with tuner.replay(((concurrent_runner, tactic),)), torch.inference_mode(): + reloaded_output = gate(hidden_states) + torch.cuda.synchronize() + torch.testing.assert_close(reloaded_output, reloaded_expected, rtol=1e-2, atol=atol) + + +def test_gated_mlp_bf16_locality_domain_end_to_end_rubin(): + import torch.nn.functional as F + + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.modules.gated_mlp import GatedMLP + + skip_if_no_locality_domain() + torch.manual_seed(2115) + AutoTuner.get().clear_cache() + + hidden_size, intermediate_size = 128, 128 + input = torch.randn(8, hidden_size, dtype=torch.bfloat16, device="cuda") + gate_weight = torch.randn(intermediate_size, hidden_size, dtype=torch.bfloat16, device="cuda") + up_weight = torch.randn_like(gate_weight) + down_weight = torch.randn(hidden_size, intermediate_size, dtype=torch.bfloat16, device="cuda") + expected = F.linear( + F.silu(F.linear(input, gate_weight)) * F.linear(input, up_weight), + down_weight, + ) + + mlp = GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + dtype=torch.bfloat16, + config=ModelConfig( + use_cute_dsl_bf16_gemm=True, + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ), + enable_locality_domain_bf16_linear=True, + ).cuda() + mlp.gate_up_proj.load_weights([{"weight": gate_weight}, {"weight": up_weight}]) + mlp.down_proj.load_weights([{"weight": down_weight}]) + mlp.gate_up_proj.post_load_weights() + mlp.down_proj.post_load_weights() + + assert mlp.gate_up_proj.partition_plan.enabled + assert mlp.down_proj.partition_plan.enabled + assert mlp.gate_up_proj.weight.numel() == 0 + assert mlp.down_proj.weight.numel() == 0 + gate_up_shards = mlp.gate_up_proj._locality_domain_weight_shards + assert gate_up_shards is not None + torch.testing.assert_close(gate_up_shards[0]["weight"], gate_weight) + torch.testing.assert_close(gate_up_shards[1]["weight"], up_weight) + + with torch.inference_mode(): + output = mlp(input) + torch.cuda.synchronize() + + torch.testing.assert_close(output, expected, rtol=2e-2, atol=2.0) + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2, + reason="Cross-device locality domain dispatch requires at least two visible CUDA devices", +) +def test_cute_dsl_bf16_gemm_locality_domain_uses_input_device_rubin(): + """The custom op must not select locality domain streams from the caller's current device.""" + skip_if_no_locality_domain() + original_device = torch.cuda.current_device() + input_device = 1 if original_device == 0 else 0 + torch.manual_seed(107) + reset_bf16_gemm_state() + + with torch.cuda.device(input_device): + act = torch.randn(1, 128, dtype=torch.bfloat16, device="cuda") + weight_0 = torch.randn(128, 128, dtype=torch.bfloat16, device="cuda") + weight_1 = torch.randn_like(weight_0) + output = torch.empty(1, 256, dtype=torch.bfloat16, device="cuda") + expected_0 = act.float() @ weight_0.t().float() + expected_1 = act.float() @ weight_1.t().float() + + assert torch.cuda.current_device() == original_device + torch.ops.trtllm.cute_dsl_bf16_gemm_locality_domain_inplace_rubin( + act, weight_0, weight_1, output + ) + torch.cuda.synchronize(input_device) + + assert torch.cuda.current_device() == original_device + atol = bf16_matmul_atol(act.shape[-1]) + torch.testing.assert_close(output[:, :128].float(), expected_0, rtol=1e-2, atol=atol) + torch.testing.assert_close(output[:, 128:].float(), expected_1, rtol=1e-2, atol=atol) + + +@pytest.mark.parametrize( + ("kernel_variant", "bm_nk"), + [("base", (2, 64, 128, 128)), ("preferred_cluster", (2, 256, 256, 128))], +) +def test_cute_dsl_bf16_bmm_locality_domain_rubin(kernel_variant, bm_nk): + skip_if_no_locality_domain() + torch.manual_seed(456) + runner = make_bf16_bmm_runner() + + batch_size, m, n, k = bm_nk + act = torch.randn(batch_size, m, k, dtype=torch.bfloat16, device="cuda") + weight_0 = torch.randn(batch_size, n, k, dtype=torch.bfloat16, device="cuda") + weight_1 = torch.randn(batch_size, n, k, dtype=torch.bfloat16, device="cuda") + output = torch.empty(batch_size, m, n, dtype=torch.bfloat16, device="cuda") + tactics = runner.get_valid_tactics([act, weight_0, output], None) + tactic = select_bf16_tactic(tactics, kernel_variant) + + expected_0 = torch.bmm(act.float(), weight_0.transpose(1, 2).float()) + expected_1 = torch.bmm(act.float(), weight_1.transpose(1, 2).float()) + runner([act, weight_0, output], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected_0, rtol=1e-2, atol=bf16_matmul_atol(k)) + + wide_output = torch.empty(batch_size, m, n * 2, dtype=torch.bfloat16, device="cuda") + run_locality_domain_composite( + "cute_dsl_bf16_bmm_locality_domain_inplace_rubin", + (act, weight_0, weight_1, wide_output), + (expected_0, expected_1), + partition_dim=2, + kernel_variant=kernel_variant, + ) + + +@pytest.mark.parametrize( + ("projection", "k", "n", "output_padding"), + [ + pytest.param("k_b_proj", 128, 256, 64, id="k-projection"), + pytest.param("v_b_proj", 512, 64, 0, id="v-projection"), + ], +) +def test_cute_dsl_bf16_bmm_locality_domain_mla_strided_cuda_graph_rubin( + projection, k, n, output_padding +): + """Match MLA absorption's transposed A and strided output views at decode M=1.""" + del projection + skip_if_no_locality_domain() + torch.manual_seed(512) + batch_size, m = 128, 1 + # MLA starts from [tokens, heads, K] and transposes tokens/heads. The + # k-projection input is also sliced out of a wider q-head tensor. + input_padding = 64 if k == 128 else 0 + input_storage = torch.randn( + m, + batch_size, + k + input_padding, + dtype=torch.bfloat16, + device="cuda", + ) + act = input_storage[..., :k].transpose(0, 1) + assert act.stride(0) == k + input_padding + assert act.stride(-1) == 1 + + weight_0 = torch.randn(batch_size, n, k, dtype=torch.bfloat16, device="cuda") + weight_1 = torch.randn_like(weight_0) + full_n = n * 2 + output_storage = torch.empty( + m, + batch_size, + full_n + output_padding, + dtype=torch.bfloat16, + device="cuda", + ) + output = output_storage[..., :full_n].transpose(0, 1) + assert output.stride(0) == full_n + output_padding + expected_0 = torch.bmm(act.float(), weight_0.transpose(1, 2).float()) + expected_1 = torch.bmm(act.float(), weight_1.transpose(1, 2).float()) + + tactics = run_locality_domain_composite( + "cute_dsl_bf16_bmm_locality_domain_inplace_rubin", + (act, weight_0, weight_1, output), + (expected_0, expected_1), + partition_dim=2, + kernel_variant="base", + ) + assert tactics + assert all(tactic[0] == "base" and tactic[3][0] == 1 for tactic in tactics) + + +def test_bf16_linear_locality_domain_end_to_end_rubin(): + import torch.nn.functional as F + + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy + from tensorrt_llm._torch.modules.linear import Linear + + skip_if_no_locality_domain() + torch.manual_seed(2027) + reset_bf16_gemm_state() + + batch_size, seq_len, in_features, out_features = 2, 8, 128, 256 + input = torch.randn( + batch_size, + seq_len, + in_features, + dtype=torch.bfloat16, + device="cuda", + ) + weight = torch.randn(out_features, in_features, dtype=torch.bfloat16, device="cuda") + bias = torch.randn(out_features, dtype=torch.bfloat16, device="cuda") + expected = F.linear(input, weight, bias) + atol = bf16_matmul_atol(in_features) + + linear = Linear( + in_features=in_features, + out_features=out_features, + bias=True, + dtype=torch.bfloat16, + use_cute_dsl_bf16_gemm=True, + enable_locality_domain_bf16_linear=True, + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ).cuda() + linear.load_weights([{"weight": weight, "bias": bias}]) + from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader + + ModelLoader._walk_transform(linear) + shards = linear._locality_domain_weight_shards + ModelLoader._walk_cache_state(linear) + linear.post_load_weights() + assert linear._locality_domain_weight_shards is shards + + assert linear.partition_plan.enabled + assert linear.partition_plan.op_kind == "bf16_linear" + assert linear._locality_domain_weight_shards is not None + assert len(linear._locality_domain_weight_shards) == 2 + assert linear.weight.numel() == 0 + localized_weight = torch.cat( + [shard["weight"] for shard in linear._locality_domain_weight_shards], dim=0 + ) + torch.testing.assert_close(localized_weight, weight) + + with torch.inference_mode(): + output = linear(input) + torch.cuda.synchronize() + torch.testing.assert_close(output, expected, rtol=1e-2, atol=atol) + + first_generation_shards = linear._locality_domain_weight_shards + linear.pre_reload_weights() + assert linear._locality_domain_weight_shards is None + assert tuple(linear.weight.shape) == tuple(weight.shape) + + reloaded_weight = torch.randn_like(weight) + reloaded_bias = torch.randn_like(bias) + reloaded_expected = F.linear(input, reloaded_weight, reloaded_bias) + linear.load_weights([{"weight": reloaded_weight, "bias": reloaded_bias}]) + linear.post_load_weights() + + assert linear._locality_domain_weight_shards is not None + assert linear._locality_domain_weight_shards is not first_generation_shards + assert len(linear._locality_domain_weight_shards) == 2 + assert linear.weight.numel() == 0 + reloaded_localized_weight = torch.cat( + [shard["weight"] for shard in linear._locality_domain_weight_shards], dim=0 + ) + torch.testing.assert_close(reloaded_localized_weight, reloaded_weight) + + with torch.inference_mode(): + reloaded_output = linear(input) + torch.cuda.synchronize() + torch.testing.assert_close(reloaded_output, reloaded_expected, rtol=1e-2, atol=atol) + + +def test_cute_dsl_bf16_bmm_autotune_all_tactics_rubin(): + torch.manual_seed(789) + runner = make_bf16_bmm_runner() + + batch_size, m, n, k = 1, 256, 256, 128 + act = torch.randn(batch_size, m, k, dtype=torch.bfloat16, device="cuda") + weight = torch.randn(batch_size, n, k, dtype=torch.bfloat16, device="cuda") + expected = torch.bmm(act.float(), weight.transpose(1, 2).float()) + atol = bf16_matmul_atol(k) + + output = torch.empty(batch_size, m, n, dtype=torch.bfloat16, device="cuda") + with autotune(skip_dynamic_tuning_buckets=True): + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin(act, weight, output) + torch.cuda.synchronize() + torch.testing.assert_close(output.float(), expected, rtol=1e-2, atol=atol) + + expected_tactics = runner.get_valid_tactics([act, weight, output], None) + assert any(t[0] == "base" and t[1] is True for t in expected_tactics) + assert any(t[0] == "preferred_cluster" and t[1] is True for t in expected_tactics) + + tuner = AutoTuner.get() + with tuner.capture() as all_tactics: + captured_output = torch.empty_like(output) + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin(act, weight, captured_output) + torch.cuda.synchronize() + torch.testing.assert_close(captured_output.float(), expected, rtol=1e-2, atol=atol) + + tested_tactics = [] + for ((captured_runner, tactic),) in all_tactics: + replay_output = torch.empty_like(output) + tested_tactics.append(tactic) + with tuner.replay(((captured_runner, tactic),)): + torch.ops.trtllm.cute_dsl_bf16_bmm_rubin(act, weight, replay_output) + torch.cuda.synchronize() + torch.testing.assert_close(replay_output.float(), expected, rtol=1e-2, atol=atol) + + assert tested_tactics == expected_tactics + + +@pytest.mark.parametrize("mnk", [(16462, 2112, 7168), (16384, 24576, 1536)]) +def test_cute_dsl_bf16_gemm_preferred_cluster_full_coverage_rubin(mnk): + """Multi-wave full-output-coverage regression for the preferred-cluster GEMM. + + The Rubin preferred-cluster mega-kernel runs a flexible cluster launch: + the hardware downgrades some preferred (4,2)=8-CTA clusters to the + fallback (2,1) shape (always happens on 216-SM parts due to GPC + fragmentation). Both cluster populations MUST schedule over the single + preferred tile partition; if the fallback-shaped clusters use a different + partition, whole cluster work-items are never visited and the + corresponding output tiles keep stale memory -- a silent accuracy bug. + + Small single-wave shapes (one persistent wave covers every tile) never + expose this. Here M is large enough to force many waves; we NaN-poison the + output so any unvisited tile shows up as NaN rows, independent of error + tolerance, and require every preferred_cluster tactic to write the full + output and match the reference. + """ + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + m, n, k = mnk + torch.manual_seed(41) + a = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") * 0.1 + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.05 + c = torch.empty(m, n, dtype=torch.bfloat16, device="cuda") + runner = cute_dsl_custom_ops.CuteDSLBf16RubinGemmRunner(use_tvm_ffi=True) + tactics = runner.get_valid_tactics([a, w, c], None) + pc = [t for t in tactics if isinstance(t, tuple) and t[0] == "preferred_cluster"] + assert pc, "no preferred_cluster tactic offered at this multi-wave shape" + + for tactic in pc: + c.fill_(float("nan")) + runner([a, w, c], tactic=tactic) + torch.cuda.synchronize() + nan_rows = int(torch.isnan(c).any(dim=1).sum().item()) + assert nan_rows == 0, f"tactic {tactic} left {nan_rows} output rows unwritten" + worst = 0.0 + for lo in range(0, m, 4096): + hi = min(m, lo + 4096) + ref = a[lo:hi].float() @ w.float().t() + worst = max(worst, (c[lo:hi].float() - ref).abs().max().item()) + assert worst < 0.5, f"tactic {tactic} wrong output: max_abs={worst:.3f}" + + +def test_cute_dsl_bf16_split_k_locality_domain_rubin(): + """Split-K runs one tactic concurrently across both real locality domain partitions.""" + skip_if_no_locality_domain() + + torch.manual_seed(99) + reset_bf16_gemm_state() + + split_k_slices = 4 + m, n, k = 256, 256, 8192 + act = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + weight_0 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + weight_1 = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + expected_0 = act.float() @ weight_0.t().float() + expected_1 = act.float() @ weight_1.t().float() + + wide_output = torch.empty(m, n * 2, dtype=torch.bfloat16, device="cuda") + run_locality_domain_composite( + "cute_dsl_bf16_gemm_locality_domain_inplace_rubin", + (act, weight_0, weight_1, wide_output), + (expected_0, expected_1), + partition_dim=1, + kernel_variant="base", + split_k_slices=split_k_slices, + capture_graph=True, + # TMA reduce-add rounds each partial sum and accumulation to BF16, + # so account for all slices as well as the reduction length. + atol=bf16_matmul_atol(k) * split_k_slices, + ) diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 4d03009ef3f4..32266434b989 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -55,6 +55,7 @@ unswizzle_sf, ) from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.models.modeling_utils import QuantAlgo def swiglu_ref(x: torch.Tensor, swiglu_limit: float = float("inf")) -> torch.Tensor: @@ -4441,7 +4442,14 @@ def create_module(enable_locality_domains: bool): get_sm_version() != 107, reason="This test is only supported on Rubin (SM 107) GPUs", ) -def test_moe_module_bf16_locality_domain_lifecycle_and_forward_rubin(): +@pytest.mark.parametrize( + ("quant_algo", "rtol", "atol"), + [ + pytest.param(None, 1e-2, 2e-2, id="bf16"), + pytest.param(QuantAlgo.NVFP4, 1e-2, 0.15, id="nvfp4"), + ], +) +def test_moe_module_locality_domain_lifecycle_and_forward_rubin(quant_algo, rtol, atol): _skip_if_no_locality_domain() from _torch.moe.quantize_utils import get_test_quant_params @@ -4473,7 +4481,7 @@ def test_moe_module_bf16_locality_domain_lifecycle_and_forward_rubin(): router_logits = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( - None, input_tensor, "CUTEDSL" + quant_algo, input_tensor, "CUTEDSL" ) quantize_util = quantize_util_cls( num_experts=num_experts, @@ -4524,19 +4532,23 @@ def create_module(enable_locality_domains: bool, post_load: bool = True): backend.w2_weight.untyped_storage().data_ptr(), ) if post_load: - moe.post_load_weights() + moe.transform_weights() + moe.cache_derived_state() moe.cuda() return moe, full_w3_w1, full_w2, source_storage_ptrs - base_moe, _, _, _ = create_module(False) - locality_domain_moe, full_w3_w1, full_w2, source_storage_ptrs = create_module(True) - locality_domain_backend = locality_domain_moe.backend + base_module, _, _, _ = create_module(False) + locality_module, full_w3_w1, full_w2, source_storage_ptrs = create_module(True) + locality_domain_backend = locality_module.backend assert locality_domain_backend._locality_domain_runtime is not None assert locality_domain_backend._locality_domain_weight_shards is not None assert hasattr(locality_domain_backend, "_cached_reserved_moe_output_memset_stream") shards = locality_domain_backend._locality_domain_weight_shards + # Repeated reader walks must preserve already-localized storage. + locality_module.cache_derived_state() + assert locality_domain_backend._locality_domain_weight_shards is shards for shard in shards: assert shard["w3_w1_weight"].untyped_storage().data_ptr() != source_storage_ptrs[0] assert shard["w2_weight"].untyped_storage().data_ptr() != source_storage_ptrs[1] @@ -4545,25 +4557,37 @@ def create_module(enable_locality_domains: bool, post_load: bool = True): assert locality_domain_backend.w3_w1_weight.numel() == 0 assert locality_domain_backend.w2_weight.numel() == 0 - # The split is a one-time transform: a repeated post-load keeps the shards. - locality_domain_moe.post_load_weights() + # A second transform must preserve the localized shards. + locality_module.transform_weights() assert locality_domain_backend._locality_domain_weight_shards is shards - # Keep the production-shape lifecycle and public ConfigurableMoE - # forward integration here. Broad accuracy, autotune, capture, and - # outer-tile replay are covered by the unified backend matrix. + # Check the public MoE forward path; backend accuracy is covered elsewhere. with torch.inference_mode(): - base_output = base_moe.forward(input_tensor, router_logits) - locality_domain_output = locality_domain_moe.forward(input_tensor, router_logits) + base_output = base_module(input_tensor, router_logits) + locality_domain_output = locality_module(input_tensor, router_logits) torch.cuda.synchronize() - torch.testing.assert_close(base_output, locality_domain_output, rtol=1e-2, atol=0.15) + torch.testing.assert_close(base_output, locality_domain_output, rtol=rtol, atol=atol) # A staged load skips transform_weights and only refreshes derived state. staged_moe, _, _, _ = create_module(True, post_load=False) with pytest.raises(NotImplementedError, match="staged loads"): staged_moe.cache_derived_state() + # Direct updates and explicit reload hooks both rebuild released weights. + for explicit_pre_reload in (False, True): + shards = locality_domain_backend._locality_domain_weight_shards + if explicit_pre_reload: + locality_module.pre_reload_weights() + locality_module.load_weights([weights]) + locality_module.transform_weights() + locality_module.cache_derived_state() + assert locality_domain_backend._locality_domain_weight_shards is not shards + with torch.inference_mode(): + reloaded_output = locality_module(input_tensor, router_logits) + torch.cuda.synchronize() + torch.testing.assert_close(reloaded_output, base_output, rtol=rtol, atol=atol) + @pytest.mark.skipif( get_sm_version() != 107, diff --git a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py index c1560de658cc..c28e617ecf76 100644 --- a/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py +++ b/tests/unittest/_torch/thop/parallel/test_dense_gemm_act_fusion.py @@ -411,8 +411,6 @@ def _spy(x, fp4_out=False): (107, False, False), ], ) -# TODO: drop this skip when the locality-domain wiring PR lands. -@pytest.mark.skip(reason="requires locality-domain wiring from a later PR in this series") def test_nvfp4_swiglu_blackwell_capability(monkeypatch, sm_version, partitioned, expected): """SM107 and locality domain shards must keep gate/up weights in the Rubin layout.""" import types @@ -428,6 +426,7 @@ def test_nvfp4_swiglu_blackwell_capability(monkeypatch, sm_version, partitioned, layer.has_bias = False layer.partition_plan = types.SimpleNamespace(enabled=partitioned) layer._weights_created = True + layer.quant_method = types.SimpleNamespace(quantizes_nvfp4_activations=True) layer.quant_config = types.SimpleNamespace( layer_quant_mode=types.SimpleNamespace(has_nvfp4=lambda: True), ) @@ -437,7 +436,11 @@ def test_nvfp4_swiglu_blackwell_capability(monkeypatch, sm_version, partitioned, mlp = GatedMLP.__new__(GatedMLP) torch.nn.Module.__init__(mlp) + mlp.split_gate_up = False mlp.gate_up_proj = layer + mlp.activation = F.silu + mlp.swiglu_alpha = None + mlp.swiglu_beta = None assert layer.can_use_cute_dsl_nvfp4_swiglu_blackwell() is expected assert mlp._can_fuse_gate_up_swiglu() is expected diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index b1efeacc3e5c..05aaca6952d4 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -21,9 +21,14 @@ import tensorrt_llm.quantization.utils.fp4_utils as fp4_utils from tensorrt_llm._torch.autotuner import autotune -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.modules.linear import Linear -from tensorrt_llm._torch.utils import model_extra_attrs +from tensorrt_llm._torch.cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE) +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy +from tensorrt_llm._torch.locality_domain_utils import is_locality_domain_enabled +from tensorrt_llm._torch.modules.linear import (Linear, WeightMode, + WeightsLoadingConfig) +from tensorrt_llm._torch.modules.swiglu import swiglu +from tensorrt_llm._torch.utils import Fp4QuantizedTensor, model_extra_attrs from tensorrt_llm._utils import get_sm_version from tensorrt_llm.math_utils import pad_up from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -929,3 +934,500 @@ def test_fp4_gemm_bias_per_backend(backend, mnk): # Tolerance is kept because the biased and unbiased calls may pick # different autotuner tactics, i.e. a different accumulation order. torch.testing.assert_close(out_fused, ref, rtol=1e-2, atol=5e-3) + + +def _skip_if_no_locality_domain(): + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + pytest.skip("CuTe DSL package with Rubin support is not available") + is_locality_domain_enabled.cache_clear() + if not is_locality_domain_enabled(): + pytest.skip( + "locality domain localization is not enabled/supported on this system" + ) + + +def _create_fp4_weights(output_size, hidden_size, dtype): + weight = torch.randn((output_size, hidden_size), dtype=dtype).cuda() + weight_sf_global = (448 * 6) / weight.abs().max().float() + weight_fp4, weight_sf_block = torch.ops.trtllm.fp4_quantize( + weight, weight_sf_global, scaling_vector_size, False) + weight_sf_block_unswizzled = ( + torch.ops.trtllm.block_scale_interleave_reverse( + weight_sf_block.cpu().view(pad_up(output_size, 128), -1))) + return weight_fp4, weight_sf_block, weight_sf_block_unswizzled, weight_sf_global + + +def _create_fp4_input(seq_len, hidden_size, dtype): + input_raw = torch.randn(seq_len, hidden_size, dtype=dtype).cuda() + input_sf_global = (448 * 6) / input_raw.abs().max().float() + input_fp4, input_sf_block = torch.ops.trtllm.fp4_quantize( + input_raw, input_sf_global, scaling_vector_size, False) + return Fp4QuantizedTensor(input_fp4, input_sf_block), input_sf_global + + +def _make_locality_domain_weight_dict(weight_fp4, + weight_sf_block_unswizzled, + input_sf_global, + weight_sf_global, + bias=None): + weight_dict = { + "input_scale": 1.0 / input_sf_global.cpu(), + "weight": weight_fp4.cpu(), + "weight_scale": weight_sf_block_unswizzled.view(torch.float8_e4m3fn), + "weight_scale_2": 1.0 / weight_sf_global.cpu(), + } + if bias is not None: + weight_dict["bias"] = bias + return [weight_dict] + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize( + "mnk", + [(1, 7168, 2112), (256, 7168, 2112)], +) +def test_fp4_linear_locality_domain_correctness(mnk, tmp_path): + _skip_if_no_locality_domain() + from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile + + seq_len, output_size, hidden_size = mnk + dtype = torch.bfloat16 + quant_config = QuantConfig(quant_algo=QuantAlgo.NVFP4) + + locality_domain_linear = Linear( + in_features=hidden_size, + out_features=output_size, + bias=False, + dtype=dtype, + quant_config=quant_config, + nvfp4_allowed_backends=["cutedsl"], + locality_domain_policy=LocalityDomainPolicy(enabled=True)) + assert locality_domain_linear.partition_plan.enabled + + weight_fp4, weight_sf, weight_sf_unswizzled, weight_sf_global = ( + _create_fp4_weights(output_size, hidden_size, dtype)) + input_tensor, input_sf_global = _create_fp4_input(seq_len, hidden_size, + dtype) + weight_dict = _make_locality_domain_weight_dict(weight_fp4, + weight_sf_unswizzled, + input_sf_global, + weight_sf_global) + + locality_domain_linear.load_weights(weight_dict) + locality_domain_linear = locality_domain_linear.cuda() + locality_domain_linear.post_load_weights() + + tuner = AutoTuner.get() + old_settings = (tuner.warmup, tuner.repeat, tuner.stream_delay_micro_secs) + # Compile each tactic before the autotuner captures its CUDA graph. + tuner.warmup = 1 + tuner.repeat = 1 + tuner.stream_delay_micro_secs = 10 + try: + tuner.clear_cache() + with torch.inference_mode(): + output_base = torch.ops.trtllm.nvfp4_gemm_cutlass( + input_tensor.fp4_tensor, + weight_fp4, + input_tensor.scaling_factor, + weight_sf, + 1.0 / (input_sf_global * weight_sf_global), + dtype, + ) + + tuner.clear_cache() + tuner.reset_statistics() + with torch.inference_mode(), autotune( + cache_path=str(tmp_path / "locality_domain.json")): + output_locality_domain = locality_domain_linear.forward( + input_tensor) + + assert seq_len in (1, 256), ( + f"tactic expectations are only defined for seq_len 1 and 256, got {seq_len}" + ) + op_name = ("trtllm::cute_dsl_nvfp4_gemm_locality_domain_inplace_rubin" + "::locality_domain_concurrent") + assert tuner.stats.tuned_op_profiled_configs.get(op_name, 0) > 0, str( + tuner.stats) + assert not tuner.stats.failed_profiling_count.get(op_name, set()) + + with tuner.capture() as tactics_capture, torch.inference_mode(): + output_locality_domain = locality_domain_linear.forward( + input_tensor) + + assert len(tactics_capture._captured_contexts) == 1 + context = tactics_capture._captured_contexts[0] + assert context["custom_op"] == op_name + concurrent_runner = context["runners"][0] + tactics = concurrent_runner.get_valid_tactics(context["inputs"], + OptimizationProfile()) + base_tactics = [tactic for tactic in tactics if tactic[0] == "base"] + mixed_tactics = [ + tactic for tactic in tactics if tactic[0] == "mixed_clusters" + ] + + if seq_len == 1: + assert {tactic[4] for tactic in base_tactics} == {True} + assert not mixed_tactics + replay_tactics = [base_tactics[0]] + else: + assert {tactic[4] for tactic in base_tactics} == {False} + assert {tactic[5] for tactic in mixed_tactics} == {True} + replay_tactics = [base_tactics[0], mixed_tactics[0]] + + for tactic in replay_tactics: + with torch.inference_mode(): + context["inputs"][-1].zero_() + concurrent_runner(context["inputs"], tactic=tactic) + torch.cuda.synchronize() + torch.testing.assert_close( + output_base, + context["inputs"][-1][:, :output_size], + rtol=1e-2, + atol=0.15, + ) + finally: + tuner.warmup, tuner.repeat, tuner.stream_delay_micro_secs = old_settings + + torch.cuda.synchronize() + assert locality_domain_linear.partition_plan.num_partitions == 2 + torch.testing.assert_close(output_base, + output_locality_domain, + rtol=1e-2, + atol=0.15) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize("input_kind", ["quantized-tensor", "tuple"]) +def test_fp4_linear_locality_domain_rank3_prequantized_input(input_kind): + _skip_if_no_locality_domain() + torch.manual_seed(0) + + batch_size, seq_len = 2, 3 + output_size, hidden_size = 192, 128 + dtype = torch.bfloat16 + common_kwargs = { + "in_features": hidden_size, + "out_features": output_size, + "bias": False, + "dtype": dtype, + "quant_config": QuantConfig(quant_algo=QuantAlgo.NVFP4), + "nvfp4_allowed_backends": ["cutlass"], + } + base_linear = Linear(**common_kwargs, + locality_domain_policy=LocalityDomainPolicy( + enabled=False)) + locality_domain_linear = Linear(**common_kwargs, + locality_domain_policy=LocalityDomainPolicy( + enabled=True)) + + input_flat, input_sf_global = _create_fp4_input(batch_size * seq_len, + hidden_size, dtype) + reference_input = Fp4QuantizedTensor( + input_flat.fp4_tensor.reshape(batch_size, seq_len, + input_flat.fp4_tensor.shape[-1]), + input_flat.scaling_factor, + input_flat.is_sf_swizzled, + ) + input_tensor = (reference_input if input_kind == "quantized-tensor" else + (reference_input.fp4_tensor, + reference_input.scaling_factor)) + weight_fp4, _, weight_sf_unswizzled, weight_sf_global = ( + _create_fp4_weights(output_size, hidden_size, dtype)) + weight_dict = _make_locality_domain_weight_dict( + weight_fp4, + weight_sf_unswizzled, + input_sf_global, + weight_sf_global, + ) + base_linear.load_weights(weight_dict) + base_linear = base_linear.cuda() + base_linear.post_load_weights() + locality_domain_linear.load_weights(weight_dict) + locality_domain_linear = locality_domain_linear.cuda() + locality_domain_linear.post_load_weights() + assert locality_domain_linear.partition_plan.enabled + + with torch.inference_mode(), autotune(): + output_base = base_linear(reference_input) + output_locality_domain = locality_domain_linear(input_tensor) + + assert output_locality_domain.shape == (batch_size, seq_len, output_size) + torch.testing.assert_close(output_locality_domain, + output_base, + rtol=1e-2, + atol=0.15) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +def test_fp4_fused_gate_up_linear_locality_domain_correctness(): + _skip_if_no_locality_domain() + torch.manual_seed(0) + + m, intermediate_size, hidden_size = 1, 1024, 2048 + output_size = 2 * intermediate_size + dtype = torch.bfloat16 + weights_loading_config = WeightsLoadingConfig( + weight_mode=WeightMode.FUSED_GATE_UP_LINEAR) + locality_domain_linear = Linear( + in_features=hidden_size, + out_features=output_size, + bias=False, + dtype=dtype, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + weights_loading_config=weights_loading_config, + fused_weight_shard_indices_mapping={ + "gate": (0, intermediate_size), + "up": (intermediate_size, intermediate_size), + }, + use_cute_dsl_blockscaling_mm=True, + nvfp4_allowed_backends=["cutlass"], + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ) + assert locality_domain_linear.partition_plan.enabled + + input_tensor, input_sf_global = _create_fp4_input(m, hidden_size, dtype) + full_weight = torch.randn(output_size, + hidden_size, + dtype=dtype, + device="cuda") + weight_sf_global = (448 * 6) / full_weight.abs().max().float() + weight_dicts = [] + quantized_weights = [] + quantized_weight_scales = [] + for weight in full_weight.chunk(2, dim=0): + weight_fp4, weight_sf = torch.ops.trtllm.fp4_quantize( + weight, weight_sf_global, scaling_vector_size, False) + weight_sf_unswizzled = (torch.ops.trtllm.block_scale_interleave_reverse( + weight_sf.cpu().view(pad_up(intermediate_size, 128), -1))) + quantized_weights.append(weight_fp4) + quantized_weight_scales.append(weight_sf) + weight_dicts.append( + _make_locality_domain_weight_dict( + weight_fp4, + weight_sf_unswizzled, + input_sf_global, + weight_sf_global, + )[0]) + + locality_domain_linear.load_weights(weight_dicts) + for attr in ( + "tmp_nvfp4_weight_scales", + "tmp_nvfp4_input_scales_list", + "tmp_nvfp4_weight_scale_2_list", + "tmp_nvfp4_pre_quant_scale", + ): + assert not hasattr(locality_domain_linear, attr) + locality_domain_linear = locality_domain_linear.cuda() + locality_domain_linear.post_load_weights() + + shards = locality_domain_linear._locality_domain_weight_shards + assert shards is not None + for shard, weight, weight_scale in zip(shards, quantized_weights, + quantized_weight_scales): + assert torch.equal(shard["weight"], weight) + assert torch.equal(shard["weight_scale"], weight_scale) + assert locality_domain_linear.weight.numel() == 0 + assert locality_domain_linear.weight_scale.numel() == 0 + + with torch.inference_mode(), autotune(): + alpha = 1.0 / (input_sf_global * weight_sf_global) + base_gate_up = torch.cat([ + torch.ops.trtllm.nvfp4_gemm_cutlass( + input_tensor.fp4_tensor, + weight, + input_tensor.scaling_factor, + weight_scale, + alpha, + dtype, + ) for weight, weight_scale in zip(quantized_weights, + quantized_weight_scales) + ], + dim=-1) + locality_domain_gate_up = locality_domain_linear(input_tensor) + base_activated = swiglu(base_gate_up) + locality_domain_activated = swiglu(locality_domain_gate_up) + + torch.cuda.synchronize() + torch.testing.assert_close(locality_domain_gate_up, + base_gate_up, + rtol=1e-2, + atol=0.15) + torch.testing.assert_close(locality_domain_activated, + base_activated, + rtol=3e-2, + atol=2.0) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize( + ("output_size", "hidden_size"), + [ + pytest.param(7168, 2112, id="aligned-production-shape"), + pytest.param(192, 128, id="padded-scale-reload"), + ], +) +def test_fp4_linear_locality_domain_weight_lifecycle_and_global_bias( + output_size, hidden_size): + _skip_if_no_locality_domain() + torch.manual_seed(0) + + m = 8 + dtype = torch.bfloat16 + linear = Linear( + in_features=hidden_size, + out_features=output_size, + bias=True, + dtype=dtype, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4), + nvfp4_allowed_backends=["cutedsl"], + locality_domain_policy=LocalityDomainPolicy(enabled=True), + ) + assert linear.partition_plan.enabled + + input_tensor, input_sf_global = _create_fp4_input(m, hidden_size, dtype) + + def make_weight_generation(): + weight_fp4, weight_sf, weight_sf_unswizzled, weight_sf_global = ( + _create_fp4_weights(output_size, hidden_size, dtype)) + bias = torch.randn(output_size, dtype=dtype) + weight_dict = _make_locality_domain_weight_dict( + weight_fp4, + weight_sf_unswizzled, + input_sf_global, + weight_sf_global, + bias, + ) + return (weight_fp4, weight_sf, weight_sf_unswizzled, weight_sf_global, + bias, weight_dict) + + def check_output(weight_fp4, weight_sf, weight_sf_global): + with torch.inference_mode(): + output_no_bias = linear.apply_linear(input_tensor, None) + output = linear(input_tensor) + reference = torch.ops.trtllm.nvfp4_gemm_cutlass( + input_tensor.fp4_tensor, + weight_fp4, + input_tensor.scaling_factor, + weight_sf, + 1.0 / (input_sf_global * weight_sf_global), + dtype, + ) + torch.testing.assert_close(output_no_bias, + reference, + rtol=1e-2, + atol=0.15) + torch.testing.assert_close(output, + output_no_bias + linear.bias, + rtol=1e-2, + atol=0.15) + + first_generation = make_weight_generation() + first_weight, first_scale, _, first_scale_global, first_bias, weight_dict = ( + first_generation) + linear.load_weights(weight_dict) + linear = linear.cuda() + full_weight = linear.weight.data.clone() + full_bias = linear.bias.data.clone() + linear.post_load_weights() + shards = linear._locality_domain_weight_shards + assert shards is not None + assert len(shards) == linear.partition_plan.num_partitions + + layout = linear.partition_plan.layout + assert layout is not None + assert layout.padded_axis_extent == full_weight.size(0) + partition_n = layout.per_partition_axis_extent(padded=True) + for shard in shards: + assert shard["weight"].shape == (partition_n, full_weight.shape[1]) + + reconstructed = torch.cat([shard["weight"] for shard in shards], dim=0) + assert torch.equal(reconstructed.cpu(), full_weight.cpu()) + assert linear.weight.numel() == 0 + assert linear.weight_scale.numel() == 0 + assert torch.equal(linear.bias.cpu(), full_bias.cpu()) + torch.testing.assert_close(full_bias.cpu(), first_bias) + assert all(set(shard) == {"weight", "weight_scale"} for shard in shards) + assert all("param" not in metadata + for metadata in linear.rebuild_tensor_metadata.values()) + check_output(first_weight, first_scale, first_scale_global) + + original_weight_shape = (output_size, hidden_size // 2) + original_scale_shape = tuple( + linear.rebuild_tensor_metadata["weight_scale"]["meta"].shape) + linear.pre_reload_weights() + assert linear._locality_domain_weight_shards is None + assert tuple(linear.weight.shape) == original_weight_shape + assert tuple(linear.weight_scale.shape) == original_scale_shape + assert linear.rebuild_tensor_metadata == {} + + second_generation = make_weight_generation() + (second_weight, second_scale, second_scale_unswizzled, second_scale_global, + second_bias, second_weight_dict) = second_generation + linear.load_weights(second_weight_dict) + linear.post_load_weights() + second_shards = linear._locality_domain_weight_shards + assert second_shards is not None + assert second_shards is not shards + assert torch.equal( + torch.cat([shard["weight"] for shard in second_shards], dim=0), + second_weight, + ) + + logical_scale_parts = [] + shard_n = output_size // len(second_shards) + for shard in second_shards: + padded_rows = pad_up(shard_n, 128) + scale = shard["weight_scale"].view(padded_rows, -1) + scale = torch.ops.trtllm.block_scale_interleave_reverse(scale) + logical_scale_parts.append(scale[:shard_n]) + reloaded_scale = torch.cat(logical_scale_parts, dim=0) + torch.testing.assert_close(reloaded_scale.cpu(), + second_scale_unswizzled[:output_size]) + assert linear.weight.numel() == 0 + assert linear.weight_scale.numel() == 0 + assert all("bias" not in shard for shard in second_shards) + torch.testing.assert_close(linear.bias.cpu(), second_bias) + check_output(second_weight, second_scale, second_scale_global) + + +@pytest.mark.skipif(get_sm_version() != 107 + or not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="Requires Rubin CuTe DSL") +def test_cute_dsl_nvfp4_non_inplace_rubin(): + torch.manual_seed(47) + x, x_scale = _create_fp4_input(8, 128, torch.bfloat16) + w, w_sf, _, w_scale = _create_fp4_weights(256, 128, torch.bfloat16) + alpha = 1.0 / (x_scale * w_scale) + args = (x.fp4_tensor, w, x.scaling_factor, w_sf, alpha, torch.bfloat16) + expected = torch.ops.trtllm.nvfp4_gemm_cutlass(*args) + op = torch.ops.trtllm.cute_dsl_nvfp4_gemm_rubin + output = op(*args) + torch.testing.assert_close(output, expected, rtol=1e-2, atol=0.15) + assert output.shape == (8, 256) + assert output.dtype == torch.bfloat16 + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = op(*args) + captured.fill_(float("nan")) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, expected, rtol=1e-2, atol=0.15) + + with pytest.raises(ValueError, match="partition_id must be -1"): + op(*args, partition_id=0) + with pytest.raises(ValueError, match="must use"): + op(*args, output_tensor=output) diff --git a/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py b/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py index a749ffc429a9..d9a7ec58253c 100644 --- a/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py +++ b/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py @@ -486,7 +486,6 @@ def test_bf16_linear_decision_table( @patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", return_value=True ) - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") @patch("tensorrt_llm._torch.cute_dsl_utils.IS_CUTLASS_DSL_AVAILABLE", True) def test_bf16_fused_gate_up_requires_canonical_mapping(self, mock_locality_domain): from tensorrt_llm._torch.modules.linear import Linear, WeightMode, WeightsLoadingConfig @@ -515,7 +514,6 @@ def test_bf16_fused_gate_up_requires_canonical_mapping(self, mock_locality_domai assert not linear.partition_plan.enabled assert "canonical [gate | up]" in linear.partition_plan.reason_if_disabled - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") def test_linear_reads_policy_from_model_extra_attrs(self): from tensorrt_llm._torch.modules.linear import Linear @@ -535,7 +533,6 @@ def test_linear_reads_policy_from_model_extra_attrs(self): @patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", return_value=True ) - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") @patch("tensorrt_llm._torch.cute_dsl_utils.IS_CUTLASS_DSL_AVAILABLE", True) def test_linear_copies_backend_list_before_locality_domain_append(self, mock_locality_domain): from tensorrt_llm._torch.modules.linear import Linear @@ -568,7 +565,6 @@ def test_linear_copies_backend_list_before_locality_domain_append(self, mock_loc @patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", return_value=True ) - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") @patch("tensorrt_llm._torch.cute_dsl_utils.IS_CUTLASS_DSL_AVAILABLE", True) @pytest.mark.parametrize("start_nvfp4", [True, False], ids=["nvfp4-to-bf16", "bf16-to-nvfp4"]) def test_linear_replans_after_quant_override(self, mock_locality_domain, start_nvfp4): @@ -603,7 +599,6 @@ def test_linear_replans_after_quant_override(self, mock_locality_domain, start_n @patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", return_value=True ) - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") @patch("tensorrt_llm._torch.cute_dsl_utils.IS_CUTLASS_DSL_AVAILABLE", True) def test_linear_replan_after_sharding_requires_reload(self, mock_locality_domain): from tensorrt_llm._torch.modules.linear import Linear @@ -631,7 +626,6 @@ def test_linear_replan_after_sharding_requires_reload(self, mock_locality_domain assert linear._locality_domain_runtime is original_runtime assert linear._locality_domain_weight_shards is shards - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") def test_model_config_exports_locality_domain_policy(self): from tensorrt_llm._torch.model_config import ModelConfig diff --git a/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py b/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py index a5b7b5480104..3cfdbf629fb5 100644 --- a/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py +++ b/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py @@ -106,7 +106,6 @@ def test_is_locality_domain_enabled_allows_rubin_when_supported(self): class TestLocalityDomainComputeTopology: """Pure mocked tests for compute topology and grid sizing.""" - @pytest.mark.skip(reason="requires _get_full_device_max_active_clusters (cute_dsl_custom_ops)") @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") def test_cluster_occupancy_cache_keeps_topology_scaling_dynamic(self, monkeypatch): from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops @@ -144,7 +143,6 @@ def get_max_active_clusters(self, cluster_size): finally: occupancy_cache.cache_clear() - @pytest.mark.skip(reason="requires _get_full_device_max_active_clusters (cute_dsl_custom_ops)") @pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, reason="cutlass-dsl is not available") def test_cluster_occupancy_cache_is_scoped_by_device_and_cluster(self, monkeypatch): from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops @@ -349,7 +347,6 @@ class TestLocalityDomainLinearRouting: "fusion_op", [AllReduceFusionOp.NONE, AllReduceFusionOp.RESIDUAL_RMS_NORM], ) - @pytest.mark.skip(reason="requires the Linear/model_config wire-up (call-site PR)") def test_shards_bypass_full_weight_paths(self, fusion_op): input_tensor = torch.ones(2, 4) local_output = torch.full((2, 3), 2.0) diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index dec6da6c5072..efe26b2e5bfb 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -14,6 +14,10 @@ methods: annotation: bool default: False status: beta + enable_locality_domains: + annotation: bool + default: False + status: prototype cp_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.CpConfig] default: null diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index a724e6e449de..2af2630ec380 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -137,6 +137,15 @@ commands: is_flag: true flags: - "--enable_chunked_prefill" + enable_locality_domains: + type: bool + default: false + status: prototype + required: false + multiple: false + is_flag: true + flags: + - "--enable_locality_domains" enable_visual_gen: type: bool default: false diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index f4bb6c6ba028..25f9711af4e6 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3304,6 +3304,33 @@ def test_serve_filters_default_values(self): assert llm_args.get("max_batch_size") == 128 assert llm_args.get("tensor_parallel_size") == 4 + def test_serve_enable_locality_domain_flag(self): + llm_args, _ = get_llm_args(model=llama_model_path, backend="pytorch") + assert llm_args.get("enable_locality_domains", False) is False + + llm_args, _ = get_llm_args(model=llama_model_path, + backend="pytorch", + enable_locality_domains=True) + + assert llm_args["enable_locality_domains"] is True + + @pytest.mark.cpu_only + @pytest.mark.parametrize("enabled", [False, True]) + @pytest.mark.parametrize("load_format", ["AUTO", "GMS"]) + def test_locality_domains_config_validation( + self, enabled: bool, load_format: Literal["AUTO", "GMS"]) -> None: + assert "enable_locality_domains" not in BaseLlmArgs.model_fields + kwargs = dict(model=llama_model_path, + enable_locality_domains=enabled, + load_format=load_format) + if enabled and load_format == "GMS": + with pytest.raises(ValidationError, + match="GMS.*enable_locality_domains"): + TorchLlmArgs(**kwargs) + else: + args = TorchLlmArgs(**kwargs) + assert args.enable_locality_domains is enabled + def test_serve_video_pruning_rate_maps_to_multimodal_config(self): llm_args, _ = get_llm_args( model=llama_model_path,