From 69c16686e347a3869025a0e8f0f6fdd933d1d9e3 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Wed, 5 Nov 2025 07:56:17 +0000 Subject: [PATCH 01/21] unify nvfp4 gemm backend Signed-off-by: Shijie Wang --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 12 + .../_torch/custom_ops/torch_custom_ops.py | 341 +++++++++++++++++- tensorrt_llm/_torch/modules/linear.py | 39 +- .../_torch/thop/parallel/test_fp4_linear.py | 6 +- 4 files changed, 363 insertions(+), 35 deletions(-) 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 01468e0fa1f7..e9c9bafc3589 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -297,6 +297,18 @@ def cute_dsl_nvfp4_gemm_blackwell( alpha: float, output_dtype: torch.dtype, ) -> torch.Tensor: + """CuteDSL-based NVFP4 GEMM optimized for Blackwell. + + .. deprecated:: + Use :func:`torch.ops.trtllm.nvfp4_gemm_unified` instead for automatic + backend selection among CUTLASS, cuBLASLt, and CuteDSL based on + performance profiling. + """ + from tensorrt_llm.logger import logger + logger.warning_once( + "cute_dsl_nvfp4_gemm_blackwell is deprecated. Use nvfp4_gemm_unified instead " + "for automatic backend selection with better performance.", + key="cute_dsl_nvfp4_gemm_blackwell_deprecated") tuner = AutoTuner.get() diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index f32a4aa27d28..fcf63bc8228e 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -8,9 +8,13 @@ import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils from tensorrt_llm import deep_gemm from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger +from tensorrt_llm.math_utils import pad_up from ..autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) +from ..cublaslt_utils import IS_CUBLASLT_AVAILABLE +from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..modules.multi_stream_utils import do_multi_stream from ..modules.swiglu import silu_and_mul_kernel from ..utils import (ActivationType, fp4_scale_infer_shape, @@ -525,7 +529,16 @@ def nvfp4_gemm_cublaslt( output_dtype: torch.dtype, to_userbuffers: bool = False, ) -> torch.Tensor: - """cuBLASLt-based NVFP4 GEMM with heuristic-based auto-tuning.""" + """cuBLASLt-based NVFP4 GEMM with heuristic-based auto-tuning. + + .. deprecated:: + Use :func:`nvfp4_gemm_unified` instead for automatic backend selection + among CUTLASS, cuBLASLt, and CuteDSL based on performance profiling. + """ + logger.warning_once( + "nvfp4_gemm_cublaslt is deprecated. Use nvfp4_gemm_unified instead " + "for automatic backend selection with better performance.", + key="nvfp4_gemm_cublaslt_deprecated") tuner = AutoTuner.get() # Use CublasLt runner with heuristic-based tuning @@ -572,7 +585,16 @@ def nvfp4_gemm( output_dtype: torch.dtype, to_userbuffers: bool = False, ) -> torch.Tensor: - """CUTLASS-based NVFP4 GEMM with auto-tuning.""" + """CUTLASS-based NVFP4 GEMM with auto-tuning. + + .. deprecated:: + Use :func:`nvfp4_gemm_unified` instead for automatic backend selection + among CUTLASS, cuBLASLt, and CuteDSL based on performance profiling. + """ + logger.warning_once( + "nvfp4_gemm is deprecated. Use nvfp4_gemm_unified instead " + "for automatic backend selection with better performance.", + key="nvfp4_gemm_deprecated") tuner = AutoTuner.get() # Use Cutlass runner with predefined configs @@ -606,6 +628,321 @@ def _( dtype=output_dtype) +# Conditional import for CuteDSL runner +if IS_CUTLASS_DSL_AVAILABLE: + from ..custom_ops.cute_dsl_custom_ops import CuteDSLNVFP4BlackwellLinear + + +class CuteDSLNVFP4Wrapper(TunableRunner): + """Wrapper to make CuteDSL runner compatible with unified nvfp4_gemm interface. + + This wrapper handles interface differences between CuteDSL and CUTLASS/cuBLASLt: + - alpha type conversion (torch.Tensor → float) + - scale tensor shape validation + - graceful degradation if inputs don't meet CuteDSL requirements + + CuteDSL is only available on Blackwell (SM 100) and requires specific scale tensor shapes. + """ + + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), ), + constraint_specs=(ConstraintSpec(2, 0, fp4_scale_infer_shape), ), + ) + + def __init__(self, to_userbuffers: bool, output_dtype: torch.dtype): + super().__init__() + self.to_userbuffers = to_userbuffers # Recorded but not used (CuteDSL doesn't support) + self.output_dtype = output_dtype + self.inner_runner = None # Lazy initialization + self._alpha_cache = None # Cache alpha value for runner initialization + + def _validate_cutedsl_inputs(self, act_fp4: torch.Tensor, + weight: torch.Tensor, act_sf: torch.Tensor, + weight_scale: torch.Tensor) -> bool: + """Check if inputs meet CuteDSL's strict requirements. + + CuteDSL requires: + - 2D tensors for act_fp4 and weight + - 1D scale tensors with specific padded shapes (can also accept 2D and flatten) + """ + if not IS_CUTLASS_DSL_AVAILABLE: + return False + + # Check tensor dimensions + if act_fp4.dim() != 2 or weight.dim() != 2: + return False + + # Scale tensors can be 1D or 2D (will be flattened in forward) + if act_sf.dim() not in [1, 2] or weight_scale.dim() not in [1, 2]: + return False + + # Calculate expected scale tensor shapes (CuteDSL specific) + m, k = act_fp4.shape[0], act_fp4.shape[1] + n = weight.shape[0] + real_k = k * 2 # fp4 is packed in uint8 + + sf_vec_size = 16 + sf_m = pad_up(m, 128) + sf_k = pad_up(real_k // sf_vec_size, 4) + sf_n = pad_up(n, 128) + + expected_act_sf_numel = sf_m * sf_k + expected_weight_scale_numel = sf_n * sf_k + + # Validate scale tensor element counts (regardless of shape) + if act_sf.numel() != expected_act_sf_numel: + logger.debug( + f"CuteDSL: act_sf element count mismatch. Expected {expected_act_sf_numel}, got {act_sf.numel()} (shape={act_sf.shape})" + ) + return False + + if weight_scale.numel() != expected_weight_scale_numel: + logger.debug( + f"CuteDSL: weight_scale element count mismatch. Expected {expected_weight_scale_numel}, got {weight_scale.numel()} (shape={weight_scale.shape})" + ) + return False + + return True + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs) -> List[Tuple]: + """Return CuteDSL tactics only if inputs are valid and SM version is 100.""" + if not IS_CUTLASS_DSL_AVAILABLE: + return [] + + # Check SM version + if get_sm_version() != 100: + return [] + + act_fp4, weight, act_sf, weight_scale, alpha = inputs + + # Validate inputs meet CuteDSL requirements + if not self._validate_cutedsl_inputs(act_fp4, weight, act_sf, + weight_scale): + return [] + + # CuteDSL requires 1D scale tensors. If they're 2D, flatten them. + if act_sf.dim() == 2: + act_sf = act_sf.reshape(-1) + if weight_scale.dim() == 2: + weight_scale = weight_scale.reshape(-1) + + # Lazy initialize inner CuteDSL runner + if self.inner_runner is None: + try: + # Convert alpha from Tensor to float + alpha_float = alpha.item() if isinstance( + alpha, torch.Tensor) else float(alpha) + self._alpha_cache = alpha_float + self.inner_runner = CuteDSLNVFP4BlackwellLinear( + alpha_float, self.output_dtype) + logger.debug("CuteDSL runner initialized successfully") + except Exception as e: + logger.debug(f"CuteDSL runner initialization failed: {e}") + return [] + + # Get valid tactics from inner CuteDSL runner + try: + return self.inner_runner.get_valid_tactics( + [act_fp4, weight, act_sf, weight_scale], profile, **kwargs) + except Exception as e: + logger.debug(f"CuteDSL get_valid_tactics failed: {e}") + return [] + + def forward( + self, + inputs: List[torch.Tensor], + tactic=-1, + ) -> torch.Tensor: + """Forward pass through CuteDSL runner.""" + # CuteDSL only needs first 4 inputs (alpha is already in the runner) + act_fp4, weight, act_sf, weight_scale, alpha = inputs + + if self.inner_runner is None: + raise RuntimeError( + "CuteDSL runner not initialized. This should not happen if get_valid_tactics was called." + ) + + # CuteDSL requires 1D scale tensors. If they're 2D, flatten them. + if act_sf.dim() == 2: + act_sf = act_sf.reshape(-1) + if weight_scale.dim() == 2: + weight_scale = weight_scale.reshape(-1) + + return self.inner_runner(inputs=[act_fp4, weight, act_sf, weight_scale], + tactic=tactic) + + +@torch.library.custom_op("trtllm::nvfp4_gemm_unified", mutates_args=()) +def nvfp4_gemm_unified( + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + to_userbuffers: bool = False, + backend: str = "auto", +) -> torch.Tensor: + """Unified NVFP4 GEMM for Blackwell (SM 100) with automatic or manual backend selection. + + This function can automatically choose the best backend or force a specific backend: + - CUTLASS: Predefined CUTLASS configurations with auto-tuning + - cuBLASLt: Heuristic-based algorithms from cuBLASLt library + - CuteDSL: Blackwell-optimized persistent kernels (when available and inputs are valid) + + The AutoTuner profiles all available backends during the first run and caches + the best choice for each input shape. Subsequent calls use the cached selection + with zero overhead. + + Args: + act_fp4: Activation tensor [m, k] in FP4 format (packed in uint8) + weight: Weight tensor [n, k] in FP4 format (packed in uint8) + act_sf: Activation scale factors + weight_scale: Weight scale factors + alpha: Scaling factor (as torch.Tensor for CUTLASS/cuBLASLt compatibility) + output_dtype: Output data type + to_userbuffers: Whether to use user buffers (CUTLASS/cuBLASLt only) + backend: Backend selection, one of: + - 'auto': AutoTuner automatically selects best backend (default) + - 'cutlass': Force use CUTLASS (FP4GemmRunner) + - 'cublaslt': Force use cuBLASLt (CublasLtFP4GemmRunner) + - 'cutedsl': Force use CuteDSL (CuteDSLNVFP4Wrapper) + + Returns: + Output tensor [m, n] with dtype=output_dtype + + Raises: + ValueError: If SM version is not 100 (Blackwell), or if backend is invalid/unavailable + """ + + # Validate Blackwell architecture + #if get_sm_version() != 100: + # raise ValueError( + # f"nvfp4_gemm requires SM 100 (Blackwell architecture), got SM {get_sm_version()}. " + # f"NVFP4 operations are only supported on Blackwell GPUs." + # ) + + # Validate backend parameter + valid_backends = ['auto', 'cutlass', 'cublaslt', 'cutedsl'] + if backend not in valid_backends: + raise ValueError( + f"Invalid backend '{backend}'. Must be one of {valid_backends}") + + # Case 1: Auto selection (default behavior) + if backend == "auto": + tuner = AutoTuner.get() + runners = [] + + # 1. Always add CUTLASS runner (most compatible) + runners.append( + FP4GemmRunner(fp4_utils.FP4GemmType.W4A4_NVFP4_NVFP4, + to_userbuffers, output_dtype)) + + # 2. Add cuBLASLt runner if available + if IS_CUBLASLT_AVAILABLE: + runners.append(CublasLtFP4GemmRunner(to_userbuffers, output_dtype)) + logger.debug("cuBLASLt runner added to nvfp4_gemm_unified") + + # 3. Add CuteDSL runner if available + # The wrapper will handle validation and graceful degradation + if IS_CUTLASS_DSL_AVAILABLE: + runners.append(CuteDSLNVFP4Wrapper(to_userbuffers, output_dtype)) + logger.debug("CuteDSL runner added to nvfp4_gemm_unified") + + # AutoTuner automatically profiles all runners and selects the best + best_runner, best_tactic = tuner.choose_one( + "trtllm::nvfp4_gemm_unified::gemm", + runners, + FP4GemmRunner. + tuning_config, # All runners use the same tuning_config + [act_fp4, weight, act_sf, weight_scale, alpha], + ) + + logger.debug( + f"nvfp4_gemm_unified selected backend: {type(best_runner).__name__}, tactic: {best_tactic}" + ) + + # Execute with the best runner + return best_runner( + inputs=[act_fp4, weight, act_sf, weight_scale, alpha], + tactic=best_tactic) + + # Case 2: Force CUTLASS + elif backend == "cutlass": + logger.debug("nvfp4_gemm_unified: Forcing CUTLASS backend") + runner = FP4GemmRunner(fp4_utils.FP4GemmType.W4A4_NVFP4_NVFP4, + to_userbuffers, output_dtype) + return runner( + inputs=[act_fp4, weight, act_sf, weight_scale, alpha], + tactic=-1 # Use default tactic + ) + + # Case 3: Force cuBLASLt + elif backend == "cublaslt": + if not IS_CUBLASLT_AVAILABLE: + raise ValueError( + "cuBLASLt backend is not available. " + "Please check cuBLASLt installation or use backend='auto'.") + logger.debug("nvfp4_gemm_unified: Forcing cuBLASLt backend") + runner = CublasLtFP4GemmRunner(to_userbuffers, output_dtype) + return runner( + inputs=[act_fp4, weight, act_sf, weight_scale, alpha], + tactic=-1 # Use default tactic + ) + + # Case 4: Force CuteDSL + elif backend == "cutedsl": + if not IS_CUTLASS_DSL_AVAILABLE: + raise ValueError( + "CuteDSL backend is not available. " + "Please check CuteDSL installation or use backend='auto'.") + + logger.debug("nvfp4_gemm_unified: Forcing CuteDSL backend") + wrapper = CuteDSLNVFP4Wrapper(to_userbuffers, output_dtype) + + # Validate inputs and get available tactics + try: + tactics = wrapper.get_valid_tactics( + [act_fp4, weight, act_sf, weight_scale, alpha], None) + except Exception as e: + raise ValueError( + f"CuteDSL backend failed input validation: {e}\n" + "Consider using backend='auto' to fall back to other backends.") + + if not tactics: + raise ValueError( + "CuteDSL backend has no valid tactics for this input shape. " + "This may be due to:\n" + " - Unsupported matrix dimensions\n" + " - Invalid scale tensor shapes\n" + " - SM version mismatch\n" + "Consider using backend='auto' to fall back to other backends.") + + # Use first available tactic + return wrapper(inputs=[act_fp4, weight, act_sf, weight_scale, alpha], + tactic=tactics[0]) + + +@nvfp4_gemm_unified.register_fake +def _( + act_fp4: torch.Tensor, + weight: torch.Tensor, + act_sf: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output_dtype: torch.dtype, + to_userbuffers: bool = False, + backend: str = "auto", +) -> torch.Tensor: + """Fake implementation for torch.compile support.""" + return act_fp4.new_empty((act_fp4.size(0), weight.size(0)), + dtype=output_dtype) + + class FP8BatchedGemmRunner(TunableRunner): runner_dict = dict() tuning_config = None diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index e24cf5f583cb..1d7bd99e4d23 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -28,8 +28,6 @@ from ..._utils import get_sm_version, is_sm_100f from ...models.modeling_utils import QuantConfig -from ..cublaslt_utils import IS_CUBLASLT_AVAILABLE -from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE from ..utils import Fp4QuantizedTensor, unswizzle_sf @@ -916,32 +914,15 @@ def apply(self, module: Linear, input: torch.Tensor, act_fp4, act_sf = torch.ops.trtllm.fp4_quantize( input, module.input_scale, module.scaling_vector_size, False) - if IS_CUTLASS_DSL_AVAILABLE and module.use_cute_dsl_nvfp4_blockscaling_mm: - output = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( - act_fp4, module.weight, act_sf, module.weight_scale, - module.scalar_alpha, module.dtype) - elif IS_CUBLASLT_AVAILABLE and module.use_cublaslt_nvfp4_blockscaling_mm: - output = torch.ops.trtllm.nvfp4_gemm_cublaslt( - act_fp4, module.weight, act_sf, module.weight_scale, - module.alpha, module.dtype) - else: - if module.enable_cuda_core and act_fp4.shape[0] <= 8: - act_sf_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( - act_sf.view((act_fp4.shape[0] + 128 - 1) // 128 * 128, -1)) - output = torch.ops.trtllm.cuda_core_nvfp4_gemm( - act_fp4, - module.weight, - scale_a=act_sf_unswizzled, - scale_b=module.weight_scale, - alpha=module.alpha, - bias=None, - out_dtype=module.dtype or input.dtype, - ) - else: - output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, + # Backend selection: 'auto' (default) | 'cutlass' | 'cublaslt' | 'cutedsl' + backend = getattr(module, 'nvfp4_backend', 'auto') + + # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL + output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4, module.weight, act_sf, module.weight_scale, - module.alpha, module.dtype) + module.alpha, module.dtype, + False, backend) # Take the dim of out_features if padded. Make sure the output is contiguous if output.shape[-1] > module.out_features: output = output[..., :module.out_features].contiguous() @@ -2012,10 +1993,9 @@ def __init__( allreduce_strategy: AllReduceStrategy = AllReduceStrategy.AUTO, force_dynamic_quantization: bool = False, use_cute_dsl_blockscaling_mm: bool = False, - use_cute_dsl_nvfp4_blockscaling_mm: bool = False, - use_cublaslt_nvfp4_blockscaling_mm: bool = False, disable_deep_gemm: bool = False, fused_weight_shard_indices_mapping: Optional[dict] = None, + nvfp4_backend: str = "auto", ): from ..distributed import AllReduce @@ -2033,10 +2013,9 @@ def __init__( self.gather_output = gather_output self.force_dynamic_quantization = force_dynamic_quantization self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm - self.use_cute_dsl_nvfp4_blockscaling_mm = use_cute_dsl_nvfp4_blockscaling_mm - self.use_cublaslt_nvfp4_blockscaling_mm = use_cublaslt_nvfp4_blockscaling_mm self.disable_deep_gemm = disable_deep_gemm self.fused_weight_shard_indices_mapping = fused_weight_shard_indices_mapping + self.nvfp4_backend = nvfp4_backend local_in_features = in_features local_out_features = out_features diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index bc78185e54e5..4efc0f87f9f8 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -116,7 +116,7 @@ def test_fp4_linear_cute_dsl(dtype, mnk): bias=False, dtype=dtype, quant_config=qc, - use_cute_dsl_nvfp4_blockscaling_mm=True) + nvfp4_backend='cutedsl') assert l_fp4.weight.dtype == fp4_utils.float4_e2m1x2 assert l_fp4.weight_scale.dtype == fp4_utils.float4_sf_dtype @@ -179,7 +179,7 @@ def fp4_linear_perf_test(dtype, SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE): bias=False, dtype=dtype, quant_config=qc, - use_cute_dsl_nvfp4_blockscaling_mm=True) + nvfp4_backend='cutedsl') assert l_fp4.weight.dtype == fp4_utils.float4_e2m1x2 assert l_fp4.weight_scale.dtype == fp4_utils.float4_sf_dtype @@ -214,7 +214,7 @@ def fp4_linear_perf_test(dtype, SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE): bias=False, dtype=dtype, quant_config=qc, - use_cute_dsl_nvfp4_blockscaling_mm=False) + nvfp4_backend='cutlass') # Use CUTLASS as reference assert l_fp4_ref.weight.dtype == fp4_utils.float4_e2m1x2 assert l_fp4_ref.weight_scale.dtype == fp4_utils.float4_sf_dtype From 29a0e1fd656146fa5872bb04e4360d9e91a8e72e Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Thu, 6 Nov 2025 06:21:42 +0000 Subject: [PATCH 02/21] update unittest Signed-off-by: Shijie Wang --- .../_torch/thop/parallel/test_fp4_linear.py | 273 ++++++++++++++++++ 1 file changed, 273 insertions(+) diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 4efc0f87f9f8..38d22f5e55ea 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -391,6 +391,279 @@ def nvfp4_gemm_perf_test( buffer_idx = buffer_idx + 1 +@skip_pre_blackwell +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("mnk", [(128, 7168, 16384), (128, 4096, 7168)]) +def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): + """Test nvfp4_gemm_unified with auto backend selection, ensuring all tactics are tested.""" + from tensorrt_llm._torch.autotuner import AutoTuner + + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + torch.manual_seed(0) + + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() + x_sf_global = (448 * 6) / x.abs().max().float() + + w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() + w_sf_global = (448 * 6) / w.abs().max().float() + w_fp4, w_sf_block = torch.ops.trtllm.fp4_quantize(w, w_sf_global, + scaling_vector_size, + False) + + # Prepare input + with torch.inference_mode(): + x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( + x, x_sf_global, scaling_vector_size, False) + alpha_ref = 1.0 / (w_sf_global * x_sf_global) + alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() + + # Reference: Use CUTLASS backend explicitly for reference output + with torch.inference_mode(): + output_ref = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutlass') + + # Test auto backend selection with autotuning + with torch.inference_mode(), autotune(): + output_auto = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='auto') + + # Verify auto mode result matches reference + torch.cuda.synchronize() + torch.testing.assert_close(output_auto, output_ref, rtol=1e-2, atol=0.15) + + # Capture all tactics using AutoTuner.capture() + with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): + output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='auto') + + # Convert tactics generator to list for counting + all_tactics_list = list(all_tactics) + + print(f"\n{'='*80}") + print( + f"Testing nvfp4_gemm_unified with M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" + ) + print(f"Total tactics found: {len(all_tactics_list)}") + print(f"{'='*80}") + + # Test each tactic individually + for idx, tactic in enumerate(all_tactics_list): + with AutoTuner.get().replay(tactic), torch.inference_mode(): + output = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='auto') + + # Verify each tactic produces correct results + torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) + # Get runner and tactic info from the captured tactic tuple + runner, tactic_value = tactic[ + 0] # First element of tuple for single context + print( + f" ✓ Tactic {idx+1}/{len(all_tactics_list)}: {runner.__class__.__name__} tactic={tactic_value} - PASSED" + ) + + print(f"{'='*80}") + print(f"All {len(all_tactics_list)} tactics verified successfully!") + print(f"{'='*80}\n") + + +@skip_pre_blackwell +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("backend", ["cutlass", "cublaslt"]) +@pytest.mark.parametrize("mnk", [(128, 4096, 7168), (256, 2048, 4096)]) +def test_nvfp4_gemm_unified_explicit_backend(dtype, backend, mnk): + """Test nvfp4_gemm_unified with explicit backend selection.""" + from tensorrt_llm._torch.cublaslt_utils import IS_CUBLASLT_AVAILABLE + + # Skip cuBLASLt test if not available + if backend == "cublaslt" and not IS_CUBLASLT_AVAILABLE: + pytest.skip("cuBLASLt FP4 GEMM not available in this build") + + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + torch.manual_seed(0) + + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() + x_sf_global = (448 * 6) / x.abs().max().float() + + w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() + w_sf_global = (448 * 6) / w.abs().max().float() + w_fp4, w_sf_block = torch.ops.trtllm.fp4_quantize(w, w_sf_global, + scaling_vector_size, + False) + + # Prepare input + with torch.inference_mode(): + x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( + x, x_sf_global, scaling_vector_size, False) + alpha_ref = 1.0 / (w_sf_global * x_sf_global) + alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() + + # Test with explicit backend + with torch.inference_mode(): + output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend=backend) + + # Reference: Use CUTLASS backend + with torch.inference_mode(): + output_ref = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutlass') + + # Compare results + torch.cuda.synchronize() + torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) + print( + f"✓ Backend '{backend}' test passed for M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" + ) + + +@pytest.mark.skipif(sys.version_info < (3, 12), + reason="cutlass-dsl 4.1.0 requires Python 3.12+") +@pytest.mark.skipif( + get_sm_version() != 100, + reason="This test is only supported in Blackwell architecture", +) +@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, + reason="cutlass-dsl is not available") +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("mnk", [(128, 7168, 16384), (256, 4096, 7168)]) +def test_nvfp4_gemm_unified_cutedsl_all_tactics(dtype, mnk): + """Test nvfp4_gemm_unified with CuteDSL backend, ensuring all tactics are tested.""" + from tensorrt_llm._torch.autotuner import AutoTuner + + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + torch.manual_seed(0) + + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() + x_sf_global = (448 * 6) / x.abs().max().float() + + w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() + w_sf_global = (448 * 6) / w.abs().max().float() + w_fp4, w_sf_block = torch.ops.trtllm.fp4_quantize(w, w_sf_global, + scaling_vector_size, + False) + + # Prepare input + with torch.inference_mode(): + x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( + x, x_sf_global, scaling_vector_size, False) + alpha_ref = 1.0 / (w_sf_global * x_sf_global) + alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() + + # Reference: Use CUTLASS backend + with torch.inference_mode(): + output_ref = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutlass') + + # Test CuteDSL backend with autotuning + with torch.inference_mode(), autotune(): + output_cutedsl = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutedsl') + + # Verify CuteDSL result matches reference + torch.cuda.synchronize() + torch.testing.assert_close(output_cutedsl, output_ref, rtol=1e-2, atol=0.15) + + # Capture all tactics using AutoTuner.capture() + with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): + output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutedsl') + + # Convert to list and filter CuteDSL tactics + all_tactics_list = list(all_tactics) + cutedsl_tactics = [ + t for t in all_tactics_list if 'CuteDSL' in t[0][0].__class__.__name__ + ] + + print(f"\n{'='*80}") + print( + f"Testing CuteDSL tactics for M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" + ) + print(f"Total CuteDSL tactics: {len(cutedsl_tactics)}") + print(f"{'='*80}") + + # Test each CuteDSL tactic + for idx, tactic in enumerate(cutedsl_tactics): + with AutoTuner.get().replay(tactic), torch.inference_mode(): + output = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutedsl') + + torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) + runner, tactic_value = tactic[0] + print( + f" ✓ CuteDSL tactic {idx+1}/{len(cutedsl_tactics)}: {runner.__class__.__name__} tactic={tactic_value} - PASSED" + ) + + print(f"{'='*80}") + print(f"All {len(cutedsl_tactics)} CuteDSL tactics verified successfully!") + print(f"{'='*80}\n") + + @pytest.mark.skipif( get_sm_version() not in [100, 103], reason="This test is only supported in Blackwell architecture", From ff6b647e0309183fb4b6948f8b3f37c51d96e599 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Thu, 6 Nov 2025 06:54:38 +0000 Subject: [PATCH 03/21] fix bug Signed-off-by: Shijie Wang --- .../_torch/custom_ops/torch_custom_ops.py | 24 ++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index fcf63bc8228e..a7a52b867b71 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -761,10 +761,28 @@ def forward( # CuteDSL only needs first 4 inputs (alpha is already in the runner) act_fp4, weight, act_sf, weight_scale, alpha = inputs + # Lazy initialization: AutoTuner skips get_valid_tactics on cache hits, + # so we need to initialize inner_runner here if it's still None if self.inner_runner is None: - raise RuntimeError( - "CuteDSL runner not initialized. This should not happen if get_valid_tactics was called." - ) + # Validate inputs before initialization + if not self._validate_cutedsl_inputs(act_fp4, weight, act_sf, + weight_scale): + raise ValueError( + "CuteDSL backend has no valid tactics for these inputs. " + "This may be due to:\n" + " - Unsupported matrix dimensions\n" + " - Invalid scale tensor shapes\n" + " - SM version mismatch\n" + "Consider using backend='auto' to fall back to other backends." + ) + + # Initialize inner runner with alpha value + alpha_float = alpha.item() if isinstance( + alpha, torch.Tensor) else float(alpha) + self._alpha_cache = alpha_float + self.inner_runner = CuteDSLNVFP4BlackwellLinear( + alpha_float, self.output_dtype) + logger.debug("CuteDSL runner lazily initialized in forward()") # CuteDSL requires 1D scale tensors. If they're 2D, flatten them. if act_sf.dim() == 2: From 30c5a9c0f5f7fa5218d5add8bf4e554232d0d68f Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Wed, 12 Nov 2025 10:07:53 +0000 Subject: [PATCH 04/21] addressing review feedback Signed-off-by: Shijie Wang --- .../_torch/custom_ops/torch_custom_ops.py | 230 +++++++++++------- .../_torch/thop/parallel/test_fp4_linear.py | 62 +++++ 2 files changed, 210 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index a7a52b867b71..4c1035275a29 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -519,6 +519,77 @@ def forward( return result +class CudaCoreNVFP4Runner(TunableRunner): + """ + CUDA Core-based NVFP4 GEMM runner on modern architectures. + + This runner is available on: + - SM >= 100 (Blackwell and newer architectures) + - M <= 8 (small batch size limitation from kernel template) + """ + + # Shared tuning config (no tactics needed, single implementation) + tuning_config = TuningConfig() + + # Minimum supported architecture: SM100 (Blackwell) + MIN_SM_VERSION = 100 + # Maximum M dimension (from cudaCoreGemmTemplateMaxM in C++ kernel) + MAX_M_DIMENSION = 8 + + def __init__(self, to_userbuffers: bool, output_dtype: torch.dtype): + super().__init__() + self.to_userbuffers = to_userbuffers + self.output_dtype = output_dtype + + def get_valid_tactics(self, inputs: List[torch.Tensor], + profile: OptimizationProfile, **kwargs) -> List[int]: + """Return [0] if architecture and shape requirements are met, otherwise [].""" + # Check architecture support at runtime + if torch.cuda.is_available(): + capability = torch.cuda.get_device_capability( + torch.device('cuda:0')) + sm_version = capability[0] * 10 + capability[1] + if sm_version < self.MIN_SM_VERSION: + return [] + else: + return [] + + # Check M dimension limitation (kernel template constraint) + act_fp4, weight, act_sf, weight_scale, alpha = inputs + m = act_fp4.shape[0] + if m > self.MAX_M_DIMENSION: + return [] + + # Single tactic (no config variations) + return [0] + + def forward( + self, + inputs: List[torch.Tensor], + tactic: int = -1, + ) -> torch.Tensor: + act_fp4, weight, act_sf, weight_scale, alpha = inputs + + # Unswizzle the activation scale factors + # act_sf is swizzled, need to reverse it for cuda_core_nvfp4_gemm + m = act_fp4.shape[0] + act_sf_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + act_sf.view((m + 128 - 1) // 128 * 128, -1)) + + # Call CUDA Core NVFP4 GEMM + result = torch.ops.trtllm.cuda_core_nvfp4_gemm( + act_fp4, + weight, + scale_a=act_sf_unswizzled, + scale_b=weight_scale, + alpha=alpha, + bias=None, + out_dtype=self.output_dtype, + to_userbuffers=self.to_userbuffers, + ) + return result + + @torch.library.custom_op("trtllm::nvfp4_gemm_cublaslt", mutates_args=()) def nvfp4_gemm_cublaslt( act_fp4: torch.Tensor, @@ -805,12 +876,13 @@ def nvfp4_gemm_unified( to_userbuffers: bool = False, backend: str = "auto", ) -> torch.Tensor: - """Unified NVFP4 GEMM for Blackwell (SM 100) with automatic or manual backend selection. + """Unified NVFP4 GEMM with automatic or manual backend selection. This function can automatically choose the best backend or force a specific backend: - CUTLASS: Predefined CUTLASS configurations with auto-tuning - cuBLASLt: Heuristic-based algorithms from cuBLASLt library - CuteDSL: Blackwell-optimized persistent kernels (when available and inputs are valid) + - CUDA Core: CUDA Core implementation on SM >= 89 (Ada+), M <= 16 (explicit selection only) The AutoTuner profiles all available backends during the first run and caches the best choice for each input shape. Subsequent calls use the cached selection @@ -829,49 +901,86 @@ def nvfp4_gemm_unified( - 'cutlass': Force use CUTLASS (FP4GemmRunner) - 'cublaslt': Force use cuBLASLt (CublasLtFP4GemmRunner) - 'cutedsl': Force use CuteDSL (CuteDSLNVFP4Wrapper) + - 'cuda_core': Force use CUDA Core (CudaCoreNVFP4Runner, requires SM >= 89, M <= 16) Returns: Output tensor [m, n] with dtype=output_dtype Raises: - ValueError: If SM version is not 100 (Blackwell), or if backend is invalid/unavailable + ValueError: If backend is invalid/unavailable """ - # Validate Blackwell architecture - #if get_sm_version() != 100: - # raise ValueError( - # f"nvfp4_gemm requires SM 100 (Blackwell architecture), got SM {get_sm_version()}. " - # f"NVFP4 operations are only supported on Blackwell GPUs." - # ) - # Validate backend parameter - valid_backends = ['auto', 'cutlass', 'cublaslt', 'cutedsl'] + valid_backends = ['auto', 'cutlass', 'cublaslt', 'cutedsl', 'cuda_core'] if backend not in valid_backends: raise ValueError( f"Invalid backend '{backend}'. Must be one of {valid_backends}") - # Case 1: Auto selection (default behavior) - if backend == "auto": - tuner = AutoTuner.get() - runners = [] + # Build list of runners based on backend parameter + runners = [] + + # CUDA Core runner can be enabled via backend='cuda_core' (not in auto mode by default) + # This avoids AutoTuner cache incompatibility issues + if backend == "cuda_core": + # Check if architecture is supported (SM >= 89) + is_cuda_core_supported = False + if torch.cuda.is_available(): + capability = torch.cuda.get_device_capability( + torch.device('cuda:0')) + sm_version = capability[0] * 10 + capability[1] + is_cuda_core_supported = sm_version >= CudaCoreNVFP4Runner.MIN_SM_VERSION + + if is_cuda_core_supported: + runners.append(CudaCoreNVFP4Runner(to_userbuffers, output_dtype)) + logger.debug("CUDA Core runner added to nvfp4_gemm_unified") + else: + raise ValueError( + f"CUDA Core backend requires SM >= {CudaCoreNVFP4Runner.MIN_SM_VERSION} (Ada or newer). " + f"Current SM version: {sm_version if torch.cuda.is_available() else 'N/A'}. " + f"Please use backend='auto' or another backend.") - # 1. Always add CUTLASS runner (most compatible) + # Add CUTLASS runner (always available) + if backend in ["auto", "cutlass"]: runners.append( FP4GemmRunner(fp4_utils.FP4GemmType.W4A4_NVFP4_NVFP4, to_userbuffers, output_dtype)) + logger.debug("CUTLASS runner added to nvfp4_gemm_unified") - # 2. Add cuBLASLt runner if available + # Add cuBLASLt runner if available + if backend in ["auto", "cublaslt"]: if IS_CUBLASLT_AVAILABLE: runners.append(CublasLtFP4GemmRunner(to_userbuffers, output_dtype)) logger.debug("cuBLASLt runner added to nvfp4_gemm_unified") + elif backend == "cublaslt": + raise ValueError( + "cuBLASLt backend is not available. " + "Please check cuBLASLt installation or use backend='auto'.") - # 3. Add CuteDSL runner if available - # The wrapper will handle validation and graceful degradation + # Add CuteDSL runner if available + if backend in ["auto", "cutedsl"]: if IS_CUTLASS_DSL_AVAILABLE: runners.append(CuteDSLNVFP4Wrapper(to_userbuffers, output_dtype)) logger.debug("CuteDSL runner added to nvfp4_gemm_unified") + elif backend == "cutedsl": + raise ValueError( + "CuteDSL backend is not available. " + "Please check CuteDSL installation or use backend='auto'.") - # AutoTuner automatically profiles all runners and selects the best + # Validate that we have at least one runner + if not runners: + raise ValueError(f"No valid runners available for backend '{backend}'. " + "This should not happen - please report this bug.") + + logger.debug( + f"nvfp4_gemm_unified: {len(runners)} runners available for backend '{backend}': " + f"{[type(r).__name__ for r in runners]}") + + # Use AutoTuner to select best runner and tactic + # - For 'auto' mode: compare across all backends, find global optimum + # - For forced backend: only one backend in list, but still find its best tactic + tuner = AutoTuner.get() + + try: best_runner, best_tactic = tuner.choose_one( "trtllm::nvfp4_gemm_unified::gemm", runners, @@ -879,70 +988,27 @@ def nvfp4_gemm_unified( tuning_config, # All runners use the same tuning_config [act_fp4, weight, act_sf, weight_scale, alpha], ) - - logger.debug( - f"nvfp4_gemm_unified selected backend: {type(best_runner).__name__}, tactic: {best_tactic}" - ) - - # Execute with the best runner - return best_runner( - inputs=[act_fp4, weight, act_sf, weight_scale, alpha], - tactic=best_tactic) - - # Case 2: Force CUTLASS - elif backend == "cutlass": - logger.debug("nvfp4_gemm_unified: Forcing CUTLASS backend") - runner = FP4GemmRunner(fp4_utils.FP4GemmType.W4A4_NVFP4_NVFP4, - to_userbuffers, output_dtype) - return runner( - inputs=[act_fp4, weight, act_sf, weight_scale, alpha], - tactic=-1 # Use default tactic + except IndexError as e: + # Provide more helpful error message + logger.error( + f"AutoTuner failed to select a runner. Available runners: " + f"{[type(r).__name__ for r in runners]}, " + f"shapes: M={act_fp4.shape[0]}, K={act_fp4.shape[1]*2}, N={weight.shape[0]}" ) - - # Case 3: Force cuBLASLt - elif backend == "cublaslt": - if not IS_CUBLASLT_AVAILABLE: - raise ValueError( - "cuBLASLt backend is not available. " - "Please check cuBLASLt installation or use backend='auto'.") - logger.debug("nvfp4_gemm_unified: Forcing cuBLASLt backend") - runner = CublasLtFP4GemmRunner(to_userbuffers, output_dtype) - return runner( - inputs=[act_fp4, weight, act_sf, weight_scale, alpha], - tactic=-1 # Use default tactic - ) - - # Case 4: Force CuteDSL - elif backend == "cutedsl": - if not IS_CUTLASS_DSL_AVAILABLE: - raise ValueError( - "CuteDSL backend is not available. " - "Please check CuteDSL installation or use backend='auto'.") - - logger.debug("nvfp4_gemm_unified: Forcing CuteDSL backend") - wrapper = CuteDSLNVFP4Wrapper(to_userbuffers, output_dtype) - - # Validate inputs and get available tactics - try: - tactics = wrapper.get_valid_tactics( - [act_fp4, weight, act_sf, weight_scale, alpha], None) - except Exception as e: - raise ValueError( - f"CuteDSL backend failed input validation: {e}\n" - "Consider using backend='auto' to fall back to other backends.") - - if not tactics: - raise ValueError( - "CuteDSL backend has no valid tactics for this input shape. " - "This may be due to:\n" - " - Unsupported matrix dimensions\n" - " - Invalid scale tensor shapes\n" - " - SM version mismatch\n" - "Consider using backend='auto' to fall back to other backends.") - - # Use first available tactic - return wrapper(inputs=[act_fp4, weight, act_sf, weight_scale, alpha], - tactic=tactics[0]) + raise RuntimeError( + f"AutoTuner failed to find a valid (runner, tactic) pair. " + f"This may indicate that all runners returned empty tactics for the given input shape. " + f"Available runners: {[type(r).__name__ for r in runners]}, " + f"Input shape: M={act_fp4.shape[0]}, K={act_fp4.shape[1]*2}, N={weight.shape[0]}" + ) from e + + logger.debug( + f"nvfp4_gemm_unified selected backend: {type(best_runner).__name__}, " + f"tactic: {best_tactic} (mode: {backend})") + + # Execute with the best runner and tactic + return best_runner(inputs=[act_fp4, weight, act_sf, weight_scale, alpha], + tactic=best_tactic) @nvfp4_gemm_unified.register_fake diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 38d22f5e55ea..6a1e8b3a895a 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -718,6 +718,68 @@ def test_fp4_linear_cublaslt(dtype, mnk): torch.testing.assert_close(output_cublaslt, output_cutlass) +@pytest.mark.skipif( + get_sm_version() < 100, + reason="CUDA Core backend requires SM >= 100 (Blackwell or newer)", +) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("mnk", [(1, 4096, 7168), (4, 7168, 16384), + (8, 2112, 7168)]) +def test_fp4_linear_cuda_core(dtype, mnk): + """Test CUDA Core NVFP4 GEMM implementation on SM >= 100 (M <= 8)""" + + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + torch.manual_seed(0) + + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() + x_sf_global = (448 * 6) / x.abs().max().float() + + w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() + w_sf_global = (448 * 6) / w.abs().max().float() + w_fp4, w_sf_block = torch.ops.trtllm.fp4_quantize(w, w_sf_global, + scaling_vector_size, + False) + + with torch.inference_mode(): + x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( + x, x_sf_global, scaling_vector_size, False) + + alpha_ref = 1.0 / (w_sf_global * x_sf_global) + alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() + + # Reference: Use CUTLASS backend + output_ref = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutlass') + + # Test CUDA Core backend + output_cuda_core = torch.ops.trtllm.nvfp4_gemm_unified( + act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cuda_core') + + # Compare results + torch.cuda.synchronize() + torch.testing.assert_close(output_cuda_core, + output_ref, + rtol=1e-2, + atol=0.15) + print( + f"✓ CUDA Core test passed for M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" + ) + + if __name__ == "__main__": # m, n, k fp4_linear_perf_test(torch.bfloat16, 128, 7168, 16384) From ba4354cfce153b860d37e6b74ed9d2706f76d7fd Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Wed, 12 Nov 2025 10:09:16 +0000 Subject: [PATCH 05/21] fix IndexError when runners list changes dynamically Signed-off-by: Shijie Wang --- tensorrt_llm/_torch/autotuner.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/autotuner.py b/tensorrt_llm/_torch/autotuner.py index ee1408b5bfd8..545066786515 100644 --- a/tensorrt_llm/_torch/autotuner.py +++ b/tensorrt_llm/_torch/autotuner.py @@ -365,11 +365,14 @@ def search_cache( Returns: A tuple containing: [is_cache_hit, runner_id, tactic, stored_profile] + runner_id is the index in the current runners list """ - for r in runners: + for idx, r in enumerate(runners): if (cache_key := self.get_cache_key(custom_op, r, input_shapes, tuning_config)) in self.cache: - return True, *self.cache[cache_key] + # Return the current index in runners list, not the cached runner_id + cached_runner_id, tactic, min_time = self.cache[cache_key] + return True, idx, tactic, min_time return False, *self.fallback_entry() From 74858c67619f971d449730d9e239a28a9e8fb976 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Thu, 13 Nov 2025 04:12:42 +0000 Subject: [PATCH 06/21] fix bug Signed-off-by: Shijie Wang --- .../_torch/custom_ops/torch_custom_ops.py | 42 ++++++++++++------- .../_torch/thop/parallel/test_fp4_linear.py | 3 +- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 4c1035275a29..d3c712871192 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -521,10 +521,10 @@ def forward( class CudaCoreNVFP4Runner(TunableRunner): """ - CUDA Core-based NVFP4 GEMM runner on modern architectures. + CUDA Core-based NVFP4 GEMM runner. This runner is available on: - - SM >= 100 (Blackwell and newer architectures) + - SM >= 100 (Blackwell) - M <= 8 (small batch size limitation from kernel template) """ @@ -882,11 +882,12 @@ def nvfp4_gemm_unified( - CUTLASS: Predefined CUTLASS configurations with auto-tuning - cuBLASLt: Heuristic-based algorithms from cuBLASLt library - CuteDSL: Blackwell-optimized persistent kernels (when available and inputs are valid) - - CUDA Core: CUDA Core implementation on SM >= 89 (Ada+), M <= 16 (explicit selection only) + - CUDA Core: CUDA Core implementation (requires SM >= 100 and M <= 8) The AutoTuner profiles all available backends during the first run and caches the best choice for each input shape. Subsequent calls use the cached selection - with zero overhead. + with zero overhead. In 'auto' mode, backends are only considered if their + requirements are met (e.g., CUDA Core only participates when SM >= 100 and M <= 8). Args: act_fp4: Activation tensor [m, k] in FP4 format (packed in uint8) @@ -901,7 +902,7 @@ def nvfp4_gemm_unified( - 'cutlass': Force use CUTLASS (FP4GemmRunner) - 'cublaslt': Force use cuBLASLt (CublasLtFP4GemmRunner) - 'cutedsl': Force use CuteDSL (CuteDSLNVFP4Wrapper) - - 'cuda_core': Force use CUDA Core (CudaCoreNVFP4Runner, requires SM >= 89, M <= 16) + - 'cuda_core': Force use CUDA Core (CudaCoreNVFP4Runner, requires SM >= 100, M <= 8) Returns: Output tensor [m, n] with dtype=output_dtype @@ -919,25 +920,34 @@ def nvfp4_gemm_unified( # Build list of runners based on backend parameter runners = [] - # CUDA Core runner can be enabled via backend='cuda_core' (not in auto mode by default) - # This avoids AutoTuner cache incompatibility issues - if backend == "cuda_core": - # Check if architecture is supported (SM >= 89) + # Add CUDA Core runner if conditions are met + # Only instantiate when both SM version and M dimension requirements are satisfied + if backend in ["auto", "cuda_core"]: is_cuda_core_supported = False + m = act_fp4.shape[0] + sm_version = None + if torch.cuda.is_available(): capability = torch.cuda.get_device_capability( torch.device('cuda:0')) sm_version = capability[0] * 10 + capability[1] - is_cuda_core_supported = sm_version >= CudaCoreNVFP4Runner.MIN_SM_VERSION + # Check both SM version and M dimension constraints + is_cuda_core_supported = ( + sm_version >= CudaCoreNVFP4Runner.MIN_SM_VERSION + and m <= CudaCoreNVFP4Runner.MAX_M_DIMENSION) if is_cuda_core_supported: runners.append(CudaCoreNVFP4Runner(to_userbuffers, output_dtype)) - logger.debug("CUDA Core runner added to nvfp4_gemm_unified") - else: - raise ValueError( - f"CUDA Core backend requires SM >= {CudaCoreNVFP4Runner.MIN_SM_VERSION} (Ada or newer). " - f"Current SM version: {sm_version if torch.cuda.is_available() else 'N/A'}. " - f"Please use backend='auto' or another backend.") + logger.debug( + f"CUDA Core runner added to nvfp4_gemm_unified (SM={sm_version}, M={m})" + ) + elif backend == "cuda_core": + # Explicitly requested but conditions not met - raise error + error_msg = f"CUDA Core backend requires SM >= {CudaCoreNVFP4Runner.MIN_SM_VERSION} and M <= {CudaCoreNVFP4Runner.MAX_M_DIMENSION}. " + error_msg += f"Current: SM={sm_version if sm_version else 'N/A'}, M={m}. " + error_msg += "Please use backend='auto' or another backend." + raise ValueError(error_msg) + # For auto mode: silently skip if conditions not met # Add CUTLASS runner (always available) if backend in ["auto", "cutlass"]: diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 6a1e8b3a895a..277041b02d88 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -39,7 +39,8 @@ def test_fp4_linear(dtype, mnk): out_features=OUTPUT_SIZE, bias=False, dtype=dtype, - quant_config=qc) + quant_config=qc, + nvfp4_backend='cutlass') # Force CUTLASS to match reference assert l_fp4.weight.dtype == fp4_utils.float4_e2m1x2 assert l_fp4.weight_scale.dtype == fp4_utils.float4_sf_dtype From 53d67f9d4face2384439bd8849da05353c1c162f Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Thu, 13 Nov 2025 16:49:07 +0000 Subject: [PATCH 07/21] Illustrate how to implement a nested tunable op without expanding runners and tactics. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- .../_torch/custom_ops/torch_custom_ops.py | 306 +++++------------- 1 file changed, 78 insertions(+), 228 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index d3c712871192..7e83c7f85ed4 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -9,7 +9,6 @@ from tensorrt_llm import deep_gemm from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger -from tensorrt_llm.math_utils import pad_up from ..autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) @@ -699,170 +698,92 @@ def _( dtype=output_dtype) -# Conditional import for CuteDSL runner -if IS_CUTLASS_DSL_AVAILABLE: - from ..custom_ops.cute_dsl_custom_ops import CuteDSLNVFP4BlackwellLinear - - -class CuteDSLNVFP4Wrapper(TunableRunner): - """Wrapper to make CuteDSL runner compatible with unified nvfp4_gemm interface. - - This wrapper handles interface differences between CuteDSL and CUTLASS/cuBLASLt: - - alpha type conversion (torch.Tensor → float) - - scale tensor shape validation - - graceful degradation if inputs don't meet CuteDSL requirements - - CuteDSL is only available on Blackwell (SM 100) and requires specific scale tensor shapes. - """ - - tuning_config = TuningConfig( - dynamic_tensor_specs=(DynamicTensorSpec( - 0, 0, get_last_power_of_2_num_tokens_buckets, - last_positive_power_of_2), ), - constraint_specs=(ConstraintSpec(2, 0, fp4_scale_infer_shape), ), - ) +class NVFP4GemmUnifiedRunner(TunableRunner): + runner_dict = dict() + op_dict = { + "cuda_core": torch.ops.trtllm.cuda_core_nvfp4_gemm, + "cutlass": torch.ops.trtllm.nvfp4_gemm, + "cublaslt": torch.ops.trtllm.nvfp4_gemm_cublaslt, + "cutedsl": torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell, + } def __init__(self, to_userbuffers: bool, output_dtype: torch.dtype): super().__init__() - self.to_userbuffers = to_userbuffers # Recorded but not used (CuteDSL doesn't support) + self.to_userbuffers = to_userbuffers self.output_dtype = output_dtype - self.inner_runner = None # Lazy initialization - self._alpha_cache = None # Cache alpha value for runner initialization - - def _validate_cutedsl_inputs(self, act_fp4: torch.Tensor, - weight: torch.Tensor, act_sf: torch.Tensor, - weight_scale: torch.Tensor) -> bool: - """Check if inputs meet CuteDSL's strict requirements. - CuteDSL requires: - - 2D tensors for act_fp4 and weight - - 1D scale tensors with specific padded shapes (can also accept 2D and flatten) - """ - if not IS_CUTLASS_DSL_AVAILABLE: - return False - - # Check tensor dimensions - if act_fp4.dim() != 2 or weight.dim() != 2: - return False - - # Scale tensors can be 1D or 2D (will be flattened in forward) - if act_sf.dim() not in [1, 2] or weight_scale.dim() not in [1, 2]: - return False - - # Calculate expected scale tensor shapes (CuteDSL specific) - m, k = act_fp4.shape[0], act_fp4.shape[1] - n = weight.shape[0] - real_k = k * 2 # fp4 is packed in uint8 - - sf_vec_size = 16 - sf_m = pad_up(m, 128) - sf_k = pad_up(real_k // sf_vec_size, 4) - sf_n = pad_up(n, 128) - - expected_act_sf_numel = sf_m * sf_k - expected_weight_scale_numel = sf_n * sf_k - - # Validate scale tensor element counts (regardless of shape) - if act_sf.numel() != expected_act_sf_numel: - logger.debug( - f"CuteDSL: act_sf element count mismatch. Expected {expected_act_sf_numel}, got {act_sf.numel()} (shape={act_sf.shape})" - ) - return False - - if weight_scale.numel() != expected_weight_scale_numel: - logger.debug( - f"CuteDSL: weight_scale element count mismatch. Expected {expected_weight_scale_numel}, got {weight_scale.numel()} (shape={weight_scale.shape})" - ) - return False - - return True - - def get_valid_tactics(self, inputs: List[torch.Tensor], + def get_valid_tactics(self, + inputs: List[torch.Tensor], profile: OptimizationProfile, + backend: str = "auto", **kwargs) -> List[Tuple]: - """Return CuteDSL tactics only if inputs are valid and SM version is 100.""" - if not IS_CUTLASS_DSL_AVAILABLE: - return [] - - # Check SM version - if get_sm_version() != 100: - return [] - + # return valid nvfp4 gemm implementations + tactics = [] act_fp4, weight, act_sf, weight_scale, alpha = inputs - # Validate inputs meet CuteDSL requirements - if not self._validate_cutedsl_inputs(act_fp4, weight, act_sf, - weight_scale): - return [] - - # CuteDSL requires 1D scale tensors. If they're 2D, flatten them. - if act_sf.dim() == 2: - act_sf = act_sf.reshape(-1) - if weight_scale.dim() == 2: - weight_scale = weight_scale.reshape(-1) - - # Lazy initialize inner CuteDSL runner - if self.inner_runner is None: - try: - # Convert alpha from Tensor to float - alpha_float = alpha.item() if isinstance( - alpha, torch.Tensor) else float(alpha) - self._alpha_cache = alpha_float - self.inner_runner = CuteDSLNVFP4BlackwellLinear( - alpha_float, self.output_dtype) - logger.debug("CuteDSL runner initialized successfully") - except Exception as e: - logger.debug(f"CuteDSL runner initialization failed: {e}") - return [] + if backend in ["auto", "cuda_core"]: + is_cuda_core_supported = False + m = act_fp4.shape[0] + sm_version = None + + if torch.cuda.is_available(): + capability = torch.cuda.get_device_capability( + torch.device('cuda:0')) + sm_version = capability[0] * 10 + capability[1] + # Check both SM version and M dimension constraints + is_cuda_core_supported = ( + sm_version >= CudaCoreNVFP4Runner.MIN_SM_VERSION + and m <= CudaCoreNVFP4Runner.MAX_M_DIMENSION) + + if is_cuda_core_supported: + tactics.append("cuda_core") + elif backend == "cuda_core": + # Explicitly requested but conditions not met - raise error + error_msg = f"CUDA Core backend requires SM >= {CudaCoreNVFP4Runner.MIN_SM_VERSION} and M <= {CudaCoreNVFP4Runner.MAX_M_DIMENSION}. " + error_msg += f"Current: SM={sm_version if sm_version else 'N/A'}, M={m}. " + error_msg += "Please use backend='auto' or another backend." + raise ValueError(error_msg) + + # Add CUTLASS runner (always available) + if backend in ["auto", "cutlass"]: + tactics.append("cutlass") + + # Add cuBLASLt runner if available + if backend in ["auto", "cublaslt"]: + if IS_CUBLASLT_AVAILABLE: + tactics.append("cublaslt") + elif backend == "cublaslt": + raise ValueError( + "cuBLASLt backend is not available. " + "Please check cuBLASLt installation or use backend='auto'.") + + # Add CuteDSL runner if available + if backend in ["auto", "cutedsl"]: + if IS_CUTLASS_DSL_AVAILABLE: + tactics.append("cutedsl") + elif backend == "cutedsl": + raise ValueError( + "CuteDSL backend is not available. " + "Please check CuteDSL installation or use backend='auto'.") - # Get valid tactics from inner CuteDSL runner - try: - return self.inner_runner.get_valid_tactics( - [act_fp4, weight, act_sf, weight_scale], profile, **kwargs) - except Exception as e: - logger.debug(f"CuteDSL get_valid_tactics failed: {e}") - return [] + return tactics def forward( self, inputs: List[torch.Tensor], - tactic=-1, + tactic: str = "cutlass", ) -> torch.Tensor: - """Forward pass through CuteDSL runner.""" - # CuteDSL only needs first 4 inputs (alpha is already in the runner) act_fp4, weight, act_sf, weight_scale, alpha = inputs - - # Lazy initialization: AutoTuner skips get_valid_tactics on cache hits, - # so we need to initialize inner_runner here if it's still None - if self.inner_runner is None: - # Validate inputs before initialization - if not self._validate_cutedsl_inputs(act_fp4, weight, act_sf, - weight_scale): - raise ValueError( - "CuteDSL backend has no valid tactics for these inputs. " - "This may be due to:\n" - " - Unsupported matrix dimensions\n" - " - Invalid scale tensor shapes\n" - " - SM version mismatch\n" - "Consider using backend='auto' to fall back to other backends." - ) - - # Initialize inner runner with alpha value - alpha_float = alpha.item() if isinstance( - alpha, torch.Tensor) else float(alpha) - self._alpha_cache = alpha_float - self.inner_runner = CuteDSLNVFP4BlackwellLinear( - alpha_float, self.output_dtype) - logger.debug("CuteDSL runner lazily initialized in forward()") - - # CuteDSL requires 1D scale tensors. If they're 2D, flatten them. - if act_sf.dim() == 2: - act_sf = act_sf.reshape(-1) - if weight_scale.dim() == 2: - weight_scale = weight_scale.reshape(-1) - - return self.inner_runner(inputs=[act_fp4, weight, act_sf, weight_scale], - tactic=tactic) + assert tactic in self.op_dict, f"Invalid tactic: {tactic}" + return self.op_dict[tactic]( + act_fp4, + weight, + act_sf, + weight_scale, + alpha, + self.output_dtype, + self.to_userbuffers, + ) @torch.library.custom_op("trtllm::nvfp4_gemm_unified", mutates_args=()) @@ -918,72 +839,7 @@ def nvfp4_gemm_unified( f"Invalid backend '{backend}'. Must be one of {valid_backends}") # Build list of runners based on backend parameter - runners = [] - - # Add CUDA Core runner if conditions are met - # Only instantiate when both SM version and M dimension requirements are satisfied - if backend in ["auto", "cuda_core"]: - is_cuda_core_supported = False - m = act_fp4.shape[0] - sm_version = None - - if torch.cuda.is_available(): - capability = torch.cuda.get_device_capability( - torch.device('cuda:0')) - sm_version = capability[0] * 10 + capability[1] - # Check both SM version and M dimension constraints - is_cuda_core_supported = ( - sm_version >= CudaCoreNVFP4Runner.MIN_SM_VERSION - and m <= CudaCoreNVFP4Runner.MAX_M_DIMENSION) - - if is_cuda_core_supported: - runners.append(CudaCoreNVFP4Runner(to_userbuffers, output_dtype)) - logger.debug( - f"CUDA Core runner added to nvfp4_gemm_unified (SM={sm_version}, M={m})" - ) - elif backend == "cuda_core": - # Explicitly requested but conditions not met - raise error - error_msg = f"CUDA Core backend requires SM >= {CudaCoreNVFP4Runner.MIN_SM_VERSION} and M <= {CudaCoreNVFP4Runner.MAX_M_DIMENSION}. " - error_msg += f"Current: SM={sm_version if sm_version else 'N/A'}, M={m}. " - error_msg += "Please use backend='auto' or another backend." - raise ValueError(error_msg) - # For auto mode: silently skip if conditions not met - - # Add CUTLASS runner (always available) - if backend in ["auto", "cutlass"]: - runners.append( - FP4GemmRunner(fp4_utils.FP4GemmType.W4A4_NVFP4_NVFP4, - to_userbuffers, output_dtype)) - logger.debug("CUTLASS runner added to nvfp4_gemm_unified") - - # Add cuBLASLt runner if available - if backend in ["auto", "cublaslt"]: - if IS_CUBLASLT_AVAILABLE: - runners.append(CublasLtFP4GemmRunner(to_userbuffers, output_dtype)) - logger.debug("cuBLASLt runner added to nvfp4_gemm_unified") - elif backend == "cublaslt": - raise ValueError( - "cuBLASLt backend is not available. " - "Please check cuBLASLt installation or use backend='auto'.") - - # Add CuteDSL runner if available - if backend in ["auto", "cutedsl"]: - if IS_CUTLASS_DSL_AVAILABLE: - runners.append(CuteDSLNVFP4Wrapper(to_userbuffers, output_dtype)) - logger.debug("CuteDSL runner added to nvfp4_gemm_unified") - elif backend == "cutedsl": - raise ValueError( - "CuteDSL backend is not available. " - "Please check CuteDSL installation or use backend='auto'.") - - # Validate that we have at least one runner - if not runners: - raise ValueError(f"No valid runners available for backend '{backend}'. " - "This should not happen - please report this bug.") - - logger.debug( - f"nvfp4_gemm_unified: {len(runners)} runners available for backend '{backend}': " - f"{[type(r).__name__ for r in runners]}") + runner = NVFP4GemmUnifiedRunner(to_userbuffers, output_dtype) # Use AutoTuner to select best runner and tactic # - For 'auto' mode: compare across all backends, find global optimum @@ -991,34 +847,28 @@ def nvfp4_gemm_unified( tuner = AutoTuner.get() try: - best_runner, best_tactic = tuner.choose_one( + _, best_tactic = tuner.choose_one( "trtllm::nvfp4_gemm_unified::gemm", - runners, + [runner], FP4GemmRunner. tuning_config, # All runners use the same tuning_config [act_fp4, weight, act_sf, weight_scale, alpha], + backend=backend, ) except IndexError as e: # Provide more helpful error message logger.error( - f"AutoTuner failed to select a runner. Available runners: " - f"{[type(r).__name__ for r in runners]}, " f"shapes: M={act_fp4.shape[0]}, K={act_fp4.shape[1]*2}, N={weight.shape[0]}" ) raise RuntimeError( f"AutoTuner failed to find a valid (runner, tactic) pair. " - f"This may indicate that all runners returned empty tactics for the given input shape. " - f"Available runners: {[type(r).__name__ for r in runners]}, " f"Input shape: M={act_fp4.shape[0]}, K={act_fp4.shape[1]*2}, N={weight.shape[0]}" ) from e - logger.debug( - f"nvfp4_gemm_unified selected backend: {type(best_runner).__name__}, " - f"tactic: {best_tactic} (mode: {backend})") - - # Execute with the best runner and tactic - return best_runner(inputs=[act_fp4, weight, act_sf, weight_scale, alpha], - tactic=best_tactic) + return runner( + inputs=[act_fp4, weight, act_sf, weight_scale, alpha], + tactic=best_tactic, + ) @nvfp4_gemm_unified.register_fake From 63ac903bbf89b8f4d90752dcbd97cc1cd49e106c Mon Sep 17 00:00:00 2001 From: Yukun He <23156053+hyukn@users.noreply.github.com> Date: Fri, 14 Nov 2025 03:01:40 +0000 Subject: [PATCH 08/21] Solve redundant profiling issues and lightly modify unit test to avoid replay issue in recursive tuning. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> --- .../_torch/custom_ops/torch_custom_ops.py | 55 +++++++---- .../_torch/thop/parallel/test_fp4_linear.py | 96 ++++++++++--------- 2 files changed, 88 insertions(+), 63 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 7e83c7f85ed4..667357bca784 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -700,12 +700,6 @@ def _( class NVFP4GemmUnifiedRunner(TunableRunner): runner_dict = dict() - op_dict = { - "cuda_core": torch.ops.trtllm.cuda_core_nvfp4_gemm, - "cutlass": torch.ops.trtllm.nvfp4_gemm, - "cublaslt": torch.ops.trtllm.nvfp4_gemm_cublaslt, - "cutedsl": torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell, - } def __init__(self, to_userbuffers: bool, output_dtype: torch.dtype): super().__init__() @@ -772,18 +766,47 @@ def forward( self, inputs: List[torch.Tensor], tactic: str = "cutlass", + **kwargs, ) -> torch.Tensor: act_fp4, weight, act_sf, weight_scale, alpha = inputs - assert tactic in self.op_dict, f"Invalid tactic: {tactic}" - return self.op_dict[tactic]( - act_fp4, - weight, - act_sf, - weight_scale, - alpha, - self.output_dtype, - self.to_userbuffers, - ) + + if tactic == "cuda_core": + # Unswizzle the activation scale factors + # act_sf is swizzled, need to reverse it for cuda_core_nvfp4_gemm + m = act_fp4.shape[0] + act_sf_unswizzled = torch.ops.trtllm.block_scale_interleave_reverse( + act_sf.view((m + 128 - 1) // 128 * 128, -1)) + + # Call CUDA Core NVFP4 GEMM + return torch.ops.trtllm.cuda_core_nvfp4_gemm( + act_fp4, + weight, + act_sf_unswizzled, + weight_scale, + alpha, + bias=None, + out_dtype=self.output_dtype, + to_userbuffers=self.to_userbuffers) + elif tactic == "cutlass": + return torch.ops.trtllm.nvfp4_gemm(act_fp4, weight, act_sf, + weight_scale, alpha, + self.output_dtype, + self.to_userbuffers) + elif tactic == "cublaslt": + return torch.ops.trtllm.nvfp4_gemm_cublaslt(act_fp4, weight, act_sf, + weight_scale, alpha, + self.output_dtype, + self.to_userbuffers) + elif tactic == "cutedsl": + return torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( + act_fp4, weight, act_sf, weight_scale, alpha, self.output_dtype) + elif tactic == -1: + return torch.ops.trtllm.nvfp4_gemm(act_fp4, weight, act_sf, + weight_scale, alpha, + self.output_dtype, + self.to_userbuffers) + else: + raise ValueError(f"Invalid tactic: {tactic}") @torch.library.custom_op("trtllm::nvfp4_gemm_unified", mutates_args=()) diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 277041b02d88..411fe14fee7f 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -397,7 +397,7 @@ def nvfp4_gemm_perf_test( @pytest.mark.parametrize("mnk", [(128, 7168, 16384), (128, 4096, 7168)]) def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): """Test nvfp4_gemm_unified with auto backend selection, ensuring all tactics are tested.""" - from tensorrt_llm._torch.autotuner import AutoTuner + from tensorrt_llm._torch.autotuner import AutoTuner, autotune SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk torch.manual_seed(0) @@ -442,56 +442,58 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): to_userbuffers=False, backend='auto') + AutoTuner.get().print_profiling_cache() + # Verify auto mode result matches reference torch.cuda.synchronize() torch.testing.assert_close(output_auto, output_ref, rtol=1e-2, atol=0.15) - # Capture all tactics using AutoTuner.capture() - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='auto') - - # Convert tactics generator to list for counting - all_tactics_list = list(all_tactics) - - print(f"\n{'='*80}") - print( - f"Testing nvfp4_gemm_unified with M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" - ) - print(f"Total tactics found: {len(all_tactics_list)}") - print(f"{'='*80}") - - # Test each tactic individually - for idx, tactic in enumerate(all_tactics_list): - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='auto') - - # Verify each tactic produces correct results - torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) - # Get runner and tactic info from the captured tactic tuple - runner, tactic_value = tactic[ - 0] # First element of tuple for single context - print( - f" ✓ Tactic {idx+1}/{len(all_tactics_list)}: {runner.__class__.__name__} tactic={tactic_value} - PASSED" - ) - - print(f"{'='*80}") - print(f"All {len(all_tactics_list)} tactics verified successfully!") - print(f"{'='*80}\n") + # # Capture all tactics using AutoTuner.capture() + # with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): + # output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, + # weight=w_fp4, + # act_sf=x_sf_block, + # weight_scale=w_sf_block, + # alpha=alpha_tensor, + # output_dtype=dtype, + # to_userbuffers=False, + # backend='auto') + + # # Convert tactics generator to list for counting + # all_tactics_list = list(all_tactics) + + # print(f"\n{'='*80}") + # print( + # f"Testing nvfp4_gemm_unified with M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" + # ) + # print(f"Total tactics found: {len(all_tactics_list)}") + # print(f"{'='*80}") + + # # Test each tactic individually + # for idx, tactic in enumerate(all_tactics_list): + # with AutoTuner.get().replay(tactic), torch.inference_mode(): + # output = torch.ops.trtllm.nvfp4_gemm_unified( + # act_fp4=x_fp4, + # weight=w_fp4, + # act_sf=x_sf_block, + # weight_scale=w_sf_block, + # alpha=alpha_tensor, + # output_dtype=dtype, + # to_userbuffers=False, + # backend='auto') + + # # Verify each tactic produces correct results + # torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) + # # Get runner and tactic info from the captured tactic tuple + # runner, tactic_value = tactic[ + # 0] # First element of tuple for single context + # print( + # f" ✓ Tactic {idx+1}/{len(all_tactics_list)}: {runner.__class__.__name__} tactic={tactic_value} - PASSED" + # ) + + # print(f"{'='*80}") + # print(f"All {len(all_tactics_list)} tactics verified successfully!") + # print(f"{'='*80}\n") @skip_pre_blackwell From c5b04ee811157ff83873738d8d15ce34c131322d Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Wed, 19 Nov 2025 03:02:38 +0000 Subject: [PATCH 09/21] Modify the alpha in cutedsl to be a device pointer and refactor code Signed-off-by: Shijie Wang --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 44 ++++++------- .../_torch/custom_ops/torch_custom_ops.py | 22 +++---- .../dense_blockscaled_gemm_persistent.py | 62 ++++++++++--------- 3 files changed, 63 insertions(+), 65 deletions(-) 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 e9c9bafc3589..687f6be99013 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -43,17 +43,20 @@ class CuteDSLNVFP4BlackwellRunner(TunableRunner): use_cold_l2_cache=True, ) - def __init__(self, alpha: float, output_dtype: torch.dtype): + def __init__(self, output_dtype: torch.dtype): super().__init__() - self.alpha = alpha - self.output_dtype = output_dtype - assert output_dtype == torch.bfloat16 if get_sm_version() not in [100, 103]: raise ValueError( f"SM version {get_sm_version()} is not supported for {self.__class__.__name__}, it only supports SM 100 and SM 103" ) + if output_dtype != torch.bfloat16: + raise ValueError( + f"CuteDSL NVFP4 only supports bfloat16 output, got {output_dtype}" + ) + self.output_dtype = output_dtype + # rewrite the hash function because the value of self.alpha doesn't affect the tactic. def unique_id(self): return (self.output_dtype, ) @@ -147,6 +150,7 @@ def forward( self, inputs: List[torch.Tensor], tactic, + **kwargs, ) -> torch.Tensor: """ Performs fp8 blockwise gemm operation using CuTe DSL. @@ -158,7 +162,6 @@ def forward( inputs[2]: Input scale tensor of shape (k//16, m), dtype: fp8. inputs[3]: Weight scale tensor of shape (n, k//16), dtype: fp8. inputs[4]: Alpha scaling factor. dtype: float32. - inputs[5]: Output dtype, expected to be torch.bfloat16. tactic: Tiling and cluster strategy, typically a tuple (mma_tiler_mn, cluster_shape_mn). Returns: @@ -176,7 +179,7 @@ def forward( False, ] - a_tensor, b_tensor, a_sf_tensor, b_sf_tensor = inputs + a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha_tensor = inputs m, k, n = a_tensor.shape[0], a_tensor.shape[1], b_tensor.shape[0] c_tensor = torch.empty(*(m, n), dtype=self.output_dtype, @@ -204,6 +207,9 @@ def forward( b_sf_tensor, cutlass.Float8E4M3FN, 16) c_ptr = self.make_cute_dsl_global_pointer(c_tensor, cutlass.BFloat16, 16) + # Create pointer to alpha on device + alpha_ptr = self.make_cute_dsl_global_pointer( + alpha_tensor, cutlass.Float32, 4) # get stream torch_stream = torch.cuda.current_stream() @@ -254,7 +260,7 @@ def forward( kernel_a_sf_ptr, kernel_b_sf_ptr, c_ptr, - self.alpha, + alpha_ptr, # Pass alpha as device pointer max_active_clusters, stream, swap_ab, @@ -277,7 +283,7 @@ def forward( kernel_a_sf_ptr, kernel_b_sf_ptr, c_ptr, - self.alpha, + alpha_ptr, # Pass alpha as device pointer stream, ) @@ -294,32 +300,26 @@ def cute_dsl_nvfp4_gemm_blackwell( weight: torch.Tensor, input_scale: torch.Tensor, weight_scale: torch.Tensor, - alpha: float, + alpha: torch.Tensor, output_dtype: torch.dtype, ) -> torch.Tensor: """CuteDSL-based NVFP4 GEMM optimized for Blackwell. - .. deprecated:: - Use :func:`torch.ops.trtllm.nvfp4_gemm_unified` instead for automatic - backend selection among CUTLASS, cuBLASLt, and CuteDSL based on - performance profiling. + Note: + This function is primarily used internally by nvfp4_gemm_unified. + Direct usage is discouraged. Consider using nvfp4_gemm_unified instead + for automatic backend selection with better performance. """ - from tensorrt_llm.logger import logger - logger.warning_once( - "cute_dsl_nvfp4_gemm_blackwell is deprecated. Use nvfp4_gemm_unified instead " - "for automatic backend selection with better performance.", - key="cute_dsl_nvfp4_gemm_blackwell_deprecated") - tuner = AutoTuner.get() - runner = CuteDSLNVFP4BlackwellRunner(alpha, output_dtype) + runner = CuteDSLNVFP4BlackwellLinear(output_dtype) inputs = [input, weight, input_scale, weight_scale] _, best_tactic = tuner.choose_one( "trtllm::cute_dsl_nvfp4_gemm_blackwell", [runner], runner.__class__.tuning_config, inputs, - ) + output = runner(inputs, tactic=best_tactic) return output @@ -329,7 +329,7 @@ def _( mat_b: torch.Tensor, input_scale: torch.Tensor, weight_scale: torch.Tensor, - alpha: float, + alpha: torch.Tensor, # Match custom op signature output_dtype: torch.dtype, ): # [m, k] diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 667357bca784..84183ac80a3b 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -601,14 +601,11 @@ def nvfp4_gemm_cublaslt( ) -> torch.Tensor: """cuBLASLt-based NVFP4 GEMM with heuristic-based auto-tuning. - .. deprecated:: - Use :func:`nvfp4_gemm_unified` instead for automatic backend selection - among CUTLASS, cuBLASLt, and CuteDSL based on performance profiling. + Note: + This function is primarily used internally by nvfp4_gemm_unified. + Direct usage is discouraged. Consider using nvfp4_gemm_unified instead + for automatic backend selection with better performance. """ - logger.warning_once( - "nvfp4_gemm_cublaslt is deprecated. Use nvfp4_gemm_unified instead " - "for automatic backend selection with better performance.", - key="nvfp4_gemm_cublaslt_deprecated") tuner = AutoTuner.get() # Use CublasLt runner with heuristic-based tuning @@ -657,14 +654,11 @@ def nvfp4_gemm( ) -> torch.Tensor: """CUTLASS-based NVFP4 GEMM with auto-tuning. - .. deprecated:: - Use :func:`nvfp4_gemm_unified` instead for automatic backend selection - among CUTLASS, cuBLASLt, and CuteDSL based on performance profiling. + Note: + This function is primarily used internally by nvfp4_gemm_unified. + Direct usage is discouraged. Consider using nvfp4_gemm_unified instead + for automatic backend selection with better performance. """ - logger.warning_once( - "nvfp4_gemm is deprecated. Use nvfp4_gemm_unified instead " - "for automatic backend selection with better performance.", - key="nvfp4_gemm_deprecated") tuner = AutoTuner.get() # Use Cutlass runner with predefined configs diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index 407e239ed1be..6e5cc636024b 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -300,7 +300,7 @@ def __call__( sfa_tensor: cute.Tensor, sfb_tensor: cute.Tensor, c_tensor: cute.Tensor, - alpha: cutlass.Float32, + alpha: cute.Pointer, # Changed from cutlass.Float32 to device pointer max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, @@ -548,34 +548,37 @@ class SharedStorage: # GPU device kernel @cute.kernel def kernel( - self, - tiled_mma: cute.TiledMma, - tiled_mma_sfb: cute.TiledMma, - tma_atom_a: cute.CopyAtom, - mA_mkl: cute.Tensor, - tma_atom_b: cute.CopyAtom, - mB_nkl: cute.Tensor, - tma_atom_sfa: cute.CopyAtom, - mSFA_mkl: cute.Tensor, - tma_atom_sfb: cute.CopyAtom, - mSFB_nkl: cute.Tensor, - tma_atom_c: Optional[cute.CopyAtom], - mC_mnl: cute.Tensor, - cluster_layout_vmnk: cute.Layout, - cluster_layout_sfb_vmnk: cute.Layout, - a_smem_layout_staged: cute.ComposedLayout, - b_smem_layout_staged: cute.ComposedLayout, - sfa_smem_layout_staged: cute.Layout, - sfb_smem_layout_staged: cute.Layout, - c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], - epi_tile: cute.Tile, - tile_sched_params: utils.PersistentTileSchedulerParams, - epilogue_op: cutlass.Constexpr, - alpha: cutlass.Float32, + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + alpha: cute. + Pointer, # Changed from cutlass.Float32 to device pointer ): """ GPU device kernel performing the Persistent batched GEMM computation. """ + alpha_value = alpha.load().to(self.c_dtype) + warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) @@ -1248,6 +1251,7 @@ def kernel( # subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) num_prev_subtiles = tile_sched.num_tiles_executed * subtile_cnt + for subtile_idx in cutlass.range(subtile_cnt): # # Load accumulator from tensor memory buffer to register @@ -1259,8 +1263,8 @@ def kernel( # Convert to C type # acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = epilogue_op( - alpha.to(self.c_dtype) * acc_vec.to(self.c_dtype)) + acc_vec = epilogue_op(alpha_value * + acc_vec.to(self.c_dtype)) tRS_rC.store(acc_vec) # @@ -1892,7 +1896,7 @@ def wrapper( a_sf_ptr: cute.Pointer, b_sf_ptr: cute.Pointer, c_ptr: cute.Pointer, - alpha: cutlass.Float32, + alpha: cute.Pointer, # Changed from cutlass.Float32 to device pointer max_active_clusters: cutlass.Constexpr, current_stream: cuda.CUstream, swap_ab: cutlass.Constexpr = False, @@ -1913,7 +1917,7 @@ def wrapper( a_sf_ptr (cute.Pointer): Pointer to the scale factor tensor for A. b_sf_ptr (cute.Pointer): Pointer to the scale factor tensor for B. c_ptr (cute.Pointer): Pointer to the C tensor. - alpha (cutlass.Float32): Scaling factor for the GEMM output. + alpha (cute.Pointer): Pointer to alpha scaling factor on device (avoids CPU-GPU sync). max_active_clusters (cutlass.Constexpr): Maximum number of active clusters. current_stream (cuda.CUstream): CUDA stream for the operation. From 9515baee36d28a7e981a1e85104d281d407bd7e1 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Fri, 21 Nov 2025 06:35:29 +0000 Subject: [PATCH 10/21] refactor and fix bug Signed-off-by: Shijie Wang --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 79 +++- .../_torch/custom_ops/torch_custom_ops.py | 58 ++- .../dense_blockscaled_gemm_persistent.py | 17 +- .../_torch/thop/parallel/test_fp4_linear.py | 354 ++++++++---------- 4 files changed, 293 insertions(+), 215 deletions(-) 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 687f6be99013..8355bb58c3c4 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -3,6 +3,8 @@ import torch +from tensorrt_llm.logger import logger + from ..._utils import get_sm_version from ...math_utils import pad_up from ..autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, @@ -77,11 +79,48 @@ def get_valid_tactics( real_k = k * 2 batch_size = 1 sf_vec_size = 16 - # m,k + + # Fixed layout for FP4: A and B are always K-major a_major = "k" - # n, k b_major = "k" + # Data types + ab_dtype = cutlass.Float4E2M1FN + c_dtype = cutlass.BFloat16 + + # Early exit: Check K dimension alignment + # For K-major layout (A and B tensors), K is the major mode (contiguous dimension). + # 16-byte alignment requirement: K must be divisible by 32 for FP4 (128 bits / 4 bits = 32) + if real_k % 32 != 0: + logger.debug( + f"CuteDSL: K={real_k} does not meet 16-byte alignment requirement " + f"(K%32={real_k%32}, expected 0). Skipping all tactics.") + return [] + + # Optimize swap_ab candidates based on M and N alignment + # swap_ab=False → C is N-major → requires N%8==0 (BF16: 128 bits / 16 bits = 8) + # swap_ab=True → C is M-major → requires M%8==0 + m_aligned = (m % 8 == 0) + n_aligned = (n % 8 == 0) + + if not m_aligned and not n_aligned: + logger.debug( + f"CuteDSL: Neither M={m} nor N={n} meets 16-byte alignment " + f"(M%8={m%8}, N%8={n%8}). No valid C layout. Skipping all tactics." + ) + return [] + + # Only test swap_ab values that satisfy alignment + swap_ab_candidates = [] + if n_aligned: + swap_ab_candidates.append(False) # N-major layout + if m_aligned: + swap_ab_candidates.append(True) # M-major layout + + logger.debug( + f"CuteDSL: M={m}(aligned={m_aligned}), N={n}(aligned={n_aligned}), K={real_k}(aligned=True). " + f"Testing swap_ab={swap_ab_candidates}") + # full shamoo mma_tiler_mn_candidates = [ (256, 128), @@ -102,7 +141,6 @@ def get_valid_tactics( (4, 2), (4, 4), ] - swap_ab_candidates = [True, False] valid_tactics = [] for swap_ab in swap_ab_candidates: @@ -118,10 +156,10 @@ def get_valid_tactics( kernel_n = n if self.__class__.kernel_class.can_implement( - cutlass.Float4E2M1FN, # ab_dtype, + ab_dtype, cutlass.Float8E4M3FN, # sf_dtype - sf_vec_size, # sf_vec_size, - cutlass.BFloat16, # c_dtype, + sf_vec_size, + c_dtype, mma_tiler_mn, cluster_shape_mn, kernel_m, @@ -135,6 +173,9 @@ def get_valid_tactics( valid_tactics.append( (mma_tiler_mn, cluster_shape_mn, swap_ab)) + logger.debug( + f"CuteDSL: Found {len(valid_tactics)} valid tactics for M={m}, N={n}, K={real_k}" + ) return valid_tactics def make_cute_dsl_global_pointer(self, tensor: torch.Tensor, dtype, @@ -193,9 +234,27 @@ def forward( sf_k = pad_up(real_k // sf_vec_size, 4) sf_n = pad_up(n, 128) - # the scaling tensor is 1D. we need to make sure it has been padded to the correct shape - assert a_sf_tensor.shape == (sf_m * sf_k, ) - assert b_sf_tensor.shape == (sf_n * sf_k, ) + # Reshape scale factors to CuteDSL's expected format + # Input format (from CUTLASS/cuBLASLt): (m*k//16,) and (n*k//16,) + # CuteDSL format: (sf_m*sf_k,) and (sf_n*sf_k,) + # Note: This is just a view change, no memory copy + expected_a_sf_size = sf_m * sf_k + expected_b_sf_size = sf_n * sf_k + + if a_sf_tensor.numel() != expected_a_sf_size: + raise ValueError( + f"CuteDSL: act scale factor size mismatch. " + f"Expected {expected_a_sf_size} (sf_m={sf_m} * sf_k={sf_k}), " + f"got {a_sf_tensor.numel()} for shape M={m}, K={real_k}") + if b_sf_tensor.numel() != expected_b_sf_size: + raise ValueError( + f"CuteDSL: weight scale factor size mismatch. " + f"Expected {expected_b_sf_size} (sf_n={sf_n} * sf_k={sf_k}), " + f"got {b_sf_tensor.numel()} for shape N={n}, K={real_k}") + + # Reshape to CuteDSL's expected format (just a view, no copy) + a_sf_tensor = a_sf_tensor.reshape(sf_m * sf_k) + b_sf_tensor = b_sf_tensor.reshape(sf_n * sf_k) a_ptr = self.make_cute_dsl_global_pointer(a_tensor, cutlass.Float4E2M1FN, 32) @@ -313,7 +372,7 @@ def cute_dsl_nvfp4_gemm_blackwell( tuner = AutoTuner.get() runner = CuteDSLNVFP4BlackwellLinear(output_dtype) - inputs = [input, weight, input_scale, weight_scale] + inputs = [input, weight, input_scale, weight_scale, alpha] _, best_tactic = tuner.choose_one( "trtllm::cute_dsl_nvfp4_gemm_blackwell", [runner], diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 84183ac80a3b..9245d7b5166e 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1,5 +1,5 @@ from functools import lru_cache -from typing import List, Mapping, Optional, Tuple +from typing import List, Mapping, Optional, Tuple, Union import torch import triton # type: ignore[import] @@ -748,7 +748,29 @@ def get_valid_tactics(self, # Add CuteDSL runner if available if backend in ["auto", "cutedsl"]: if IS_CUTLASS_DSL_AVAILABLE: - tactics.append("cutedsl") + # Check if CuteDSL actually supports the current shape + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ + CuteDSLNVFP4BlackwellLinear + cutedsl_runner = CuteDSLNVFP4BlackwellLinear(self.output_dtype) + cutedsl_tactics = cutedsl_runner.get_valid_tactics( + inputs, profile) + + if cutedsl_tactics: + # CuteDSL supports this shape + tactics.append("cutedsl") + elif backend == "cutedsl": + # Explicitly requested CuteDSL but it doesn't support this shape + m, n, k = inputs[0].shape[0], inputs[1].shape[ + 0], inputs[0].shape[1] * 2 + raise ValueError( + f"CuteDSL backend does not support the current shape:\n" + f" M={m}, N={n}, K={k}\n" + f"CuteDSL requires 16-byte alignment for major (contiguous) dimensions:\n" + f" - K must be divisible by 32 (FP4 K-major layout): K%32={'0✓' if k % 32 == 0 else str(k%32)+'✗'}\n" + f" - Or the combination of (M, N, K, tiling, cluster shape) is not supported\n" + f"Please use backend='auto' to automatically select a compatible backend." + ) + # else: backend='auto' and CuteDSL doesn't support → silently skip elif backend == "cutedsl": raise ValueError( "CuteDSL backend is not available. " @@ -759,11 +781,40 @@ def get_valid_tactics(self, def forward( self, inputs: List[torch.Tensor], - tactic: str = "cutlass", + tactic: Union[ + str, int] = "cutlass", # str: backend name, or int: -1 for fallback **kwargs, ) -> torch.Tensor: act_fp4, weight, act_sf, weight_scale, alpha = inputs + # Check if a specific backend was requested + requested_backend = kwargs.get('backend', 'auto') + + # If a specific backend was requested (not 'auto') and we're using fallback tactic + # This can happen on cache miss, where AutoTuner uses tactic=-1 as default + if requested_backend != 'auto' and requested_backend != tactic and tactic == -1: + # User explicitly requested a backend, but we're falling back to default + # This might happen on cache miss. We should validate the requested backend supports this shape. + + # Get valid tactics for the requested backend + from tensorrt_llm._torch.autotuner import OptimizationProfile + valid_tactics = self.get_valid_tactics(inputs, + OptimizationProfile(), + backend=requested_backend) + + if not valid_tactics or requested_backend not in valid_tactics: + # Requested backend doesn't support this shape + m, n, k = inputs[0].shape[0], inputs[1].shape[ + 0], inputs[0].shape[1] * 2 + raise ValueError( + f"Backend '{requested_backend}' was explicitly requested but does not support the current shape:\n" + f" M={m}, N={n}, K={k}\n" + f"Please use backend='auto' to automatically select a compatible backend." + ) + + # Backend supports it, use the requested backend instead of fallback + tactic = requested_backend + if tactic == "cuda_core": # Unswizzle the activation scale factors # act_sf is swizzled, need to reverse it for cuda_core_nvfp4_gemm @@ -885,6 +936,7 @@ def nvfp4_gemm_unified( return runner( inputs=[act_fp4, weight, act_sf, weight_scale, alpha], tactic=best_tactic, + backend=backend, ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index 6e5cc636024b..b27b096a820e 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -300,7 +300,7 @@ def __call__( sfa_tensor: cute.Tensor, sfb_tensor: cute.Tensor, c_tensor: cute.Tensor, - alpha: cute.Pointer, # Changed from cutlass.Float32 to device pointer + alpha: cute.Tensor, # Single-element tensor containing alpha value max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, @@ -571,13 +571,12 @@ def kernel( epi_tile: cute.Tile, tile_sched_params: utils.PersistentTileSchedulerParams, epilogue_op: cutlass.Constexpr, - alpha: cute. - Pointer, # Changed from cutlass.Float32 to device pointer + alpha: cute.Tensor, # Single-element tensor containing alpha value ): """ GPU device kernel performing the Persistent batched GEMM computation. """ - alpha_value = alpha.load().to(self.c_dtype) + alpha_value = alpha[0].to(self.c_dtype) warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) @@ -1896,7 +1895,8 @@ def wrapper( a_sf_ptr: cute.Pointer, b_sf_ptr: cute.Pointer, c_ptr: cute.Pointer, - alpha: cute.Pointer, # Changed from cutlass.Float32 to device pointer + alpha: cute. + Pointer, # Device pointer to alpha, will be converted to Tensor max_active_clusters: cutlass.Constexpr, current_stream: cuda.CUstream, swap_ab: cutlass.Constexpr = False, @@ -1917,7 +1917,7 @@ def wrapper( a_sf_ptr (cute.Pointer): Pointer to the scale factor tensor for A. b_sf_ptr (cute.Pointer): Pointer to the scale factor tensor for B. c_ptr (cute.Pointer): Pointer to the C tensor. - alpha (cute.Pointer): Pointer to alpha scaling factor on device (avoids CPU-GPU sync). + alpha (cute.Pointer): Device pointer to alpha scaling factor (converted to Tensor internally). max_active_clusters (cutlass.Constexpr): Maximum number of active clusters. current_stream (cuda.CUstream): CUDA stream for the operation. @@ -1962,8 +1962,11 @@ def wrapper( (32, 4, sf_n, 4, sf_k, l), order=(2, 1, 4, 0, 3, 5), )) + alpha_tensor = cute.make_tensor(alpha, + layout=cute.make_ordered_layout( + (1, ), order=(0, ))) - self(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, alpha, + self(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, alpha_tensor, max_active_clusters, current_stream, epilogue_op) diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 411fe14fee7f..7f1a4ccb581f 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -394,12 +394,51 @@ def nvfp4_gemm_perf_test( @skip_pre_blackwell @pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("mnk", [(128, 7168, 16384), (128, 4096, 7168)]) +@pytest.mark.parametrize( + "mnk", + [ + # Small batch sizes (M <= 16) - test small M handling + (1, 4096, 4096, "Batch=1, Square 4K"), + (4, 4096, 4096, "Batch=4, Square 4K"), + (16, 4096, 4096, "Batch=16, Square 4K"), + + # Odd M values + (3, 4096, 4096, "Odd M: M=3"), + (7, 4096, 4096, "Odd M: M=7"), + (9, 4096, 4096, "Odd M: M=9"), + + # Medium batch sizes - common inference scenarios + (128, 4096, 4096, "Batch=128, Square 4K"), + (128, 7168, 16384, "Batch=128, Large K/N"), + (128, 4096, 7168, "Batch=128, Asymmetric"), + + # Large batch sizes - training scenarios + (512, 4096, 4096, "Batch=512, Square 4K"), + (1024, 4096, 4096, "Batch=1024, Square 4K"), + + # Very large batch - maximum performance + (2048, 4096, 4096, "Batch=2048, Square 4K"), + (4096, 4096, 4096, "Batch=4096, Square 4K"), + + # Large K and N - test memory bandwidth + (128, 8192, 8192, "Batch=128, Square 8K"), + (256, 16384, 16384, "Batch=256, Square 16K"), + + # Size asymmetry tests + (1024, 128, 4096, "Wide M: M >> N"), + (128, 16384, 128, "Wide N: N >> K"), + ]) def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): """Test nvfp4_gemm_unified with auto backend selection, ensuring all tactics are tested.""" from tensorrt_llm._torch.autotuner import AutoTuner, autotune + from tensorrt_llm._torch.cublaslt_utils import IS_CUBLASLT_AVAILABLE - SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + # Unpack mnk with optional description + if len(mnk) == 4: + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE, desc = mnk + else: + SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk + desc = f"M={SEQ_LEN}, K={HIDDEN_SIZE}, N={OUTPUT_SIZE}" torch.manual_seed(0) x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() @@ -448,87 +487,18 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): torch.cuda.synchronize() torch.testing.assert_close(output_auto, output_ref, rtol=1e-2, atol=0.15) - # # Capture all tactics using AutoTuner.capture() - # with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - # output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, - # weight=w_fp4, - # act_sf=x_sf_block, - # weight_scale=w_sf_block, - # alpha=alpha_tensor, - # output_dtype=dtype, - # to_userbuffers=False, - # backend='auto') - - # # Convert tactics generator to list for counting - # all_tactics_list = list(all_tactics) - - # print(f"\n{'='*80}") - # print( - # f"Testing nvfp4_gemm_unified with M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" - # ) - # print(f"Total tactics found: {len(all_tactics_list)}") - # print(f"{'='*80}") - - # # Test each tactic individually - # for idx, tactic in enumerate(all_tactics_list): - # with AutoTuner.get().replay(tactic), torch.inference_mode(): - # output = torch.ops.trtllm.nvfp4_gemm_unified( - # act_fp4=x_fp4, - # weight=w_fp4, - # act_sf=x_sf_block, - # weight_scale=w_sf_block, - # alpha=alpha_tensor, - # output_dtype=dtype, - # to_userbuffers=False, - # backend='auto') - - # # Verify each tactic produces correct results - # torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) - # # Get runner and tactic info from the captured tactic tuple - # runner, tactic_value = tactic[ - # 0] # First element of tuple for single context - # print( - # f" ✓ Tactic {idx+1}/{len(all_tactics_list)}: {runner.__class__.__name__} tactic={tactic_value} - PASSED" - # ) - - # print(f"{'='*80}") - # print(f"All {len(all_tactics_list)} tactics verified successfully!") - # print(f"{'='*80}\n") - - -@skip_pre_blackwell -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("backend", ["cutlass", "cublaslt"]) -@pytest.mark.parametrize("mnk", [(128, 4096, 7168), (256, 2048, 4096)]) -def test_nvfp4_gemm_unified_explicit_backend(dtype, backend, mnk): - """Test nvfp4_gemm_unified with explicit backend selection.""" - from tensorrt_llm._torch.cublaslt_utils import IS_CUBLASLT_AVAILABLE - - # Skip cuBLASLt test if not available - if backend == "cublaslt" and not IS_CUBLASLT_AVAILABLE: - pytest.skip("cuBLASLt FP4 GEMM not available in this build") - - SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk - torch.manual_seed(0) - - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() - x_sf_global = (448 * 6) / x.abs().max().float() - - w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() - w_sf_global = (448 * 6) / w.abs().max().float() - w_fp4, w_sf_block = torch.ops.trtllm.fp4_quantize(w, w_sf_global, - scaling_vector_size, - False) + # Test all combinations of outer layer (backend selection) and inner layer (backend tactics) + # Outer layer: nvfp4_gemm_unified selects backend + # Inner layer: each backend has its own tactics + from collections import defaultdict - # Prepare input - with torch.inference_mode(): - x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( - x, x_sf_global, scaling_vector_size, False) - alpha_ref = 1.0 / (w_sf_global * x_sf_global) - alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() + print(f"\n{'='*80}") + print(f"Testing nvfp4_gemm_unified (2-layer tactics): {desc}") + print(f"Shape: M={SEQ_LEN}, K={HIDDEN_SIZE}, N={OUTPUT_SIZE}") + print(f"{'='*80}") - # Test with explicit backend - with torch.inference_mode(): + print(f"\n[Outer Layer] Capturing backend selection tactics...") + with AutoTuner.get().capture() as outer_capture, torch.inference_mode(): output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, weight=w_fp4, act_sf=x_sf_block, @@ -536,117 +506,99 @@ def test_nvfp4_gemm_unified_explicit_backend(dtype, backend, mnk): alpha=alpha_tensor, output_dtype=dtype, to_userbuffers=False, - backend=backend) - - # Reference: Use CUTLASS backend - with torch.inference_mode(): - output_ref = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cutlass') - - # Compare results - torch.cuda.synchronize() - torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) - print( - f"✓ Backend '{backend}' test passed for M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" - ) - - -@pytest.mark.skipif(sys.version_info < (3, 12), - reason="cutlass-dsl 4.1.0 requires Python 3.12+") -@pytest.mark.skipif( - get_sm_version() != 100, - reason="This test is only supported in Blackwell architecture", -) -@pytest.mark.skipif(not IS_CUTLASS_DSL_AVAILABLE, - reason="cutlass-dsl is not available") -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("mnk", [(128, 7168, 16384), (256, 4096, 7168)]) -def test_nvfp4_gemm_unified_cutedsl_all_tactics(dtype, mnk): - """Test nvfp4_gemm_unified with CuteDSL backend, ensuring all tactics are tested.""" - from tensorrt_llm._torch.autotuner import AutoTuner - - SEQ_LEN, OUTPUT_SIZE, HIDDEN_SIZE = mnk - torch.manual_seed(0) - - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype).cuda() - x_sf_global = (448 * 6) / x.abs().max().float() - - w = torch.randn((OUTPUT_SIZE, HIDDEN_SIZE), dtype=dtype).cuda() - w_sf_global = (448 * 6) / w.abs().max().float() - w_fp4, w_sf_block = torch.ops.trtllm.fp4_quantize(w, w_sf_global, - scaling_vector_size, - False) - - # Prepare input - with torch.inference_mode(): - x_fp4, x_sf_block = torch.ops.trtllm.fp4_quantize( - x, x_sf_global, scaling_vector_size, False) - alpha_ref = 1.0 / (w_sf_global * x_sf_global) - alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() - - # Reference: Use CUTLASS backend - with torch.inference_mode(): - output_ref = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cutlass') + backend='auto') + + outer_tactics_list = list(outer_capture) + print(f" Found {len(outer_tactics_list)} outer layer tactics (backends)") + + # Parse outer tactics to get backend names + backend_map = {} + for outer_tactic in outer_tactics_list: + outer_runner, backend_name = outer_tactic[0] + backend_map[backend_name] = outer_tactic + print(f" - Backend: {backend_name}") + + print(f"\n[Inner Layer] Testing tactics for each backend...") + + # All backends have independent APIs, but cuda_core needs special handling, because it requires unswizzled scale factors + backend_apis = {} + if IS_CUTLASS_DSL_AVAILABLE: + if 'cutlass' in backend_map: + backend_apis['cutlass'] = torch.ops.trtllm.nvfp4_gemm + if IS_CUBLASLT_AVAILABLE: + if 'cublaslt' in backend_map: + backend_apis['cublaslt'] = torch.ops.trtllm.nvfp4_gemm_cublaslt + if IS_CUTLASS_DSL_AVAILABLE: + if 'cutedsl' in backend_map: + backend_apis[ + 'cutedsl'] = torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell + + # cuda_core needs special handling (different parameters, single tactic) + test_cuda_core = 'cuda_core' in backend_map + + # Step 3: For each backend, capture and immediately test all tactics + # Must test immediately after capture to avoid _last_capture being overwritten + tactics_by_backend = defaultdict(list) + total_tactics_tested = 0 + + for backend_name, backend_api in backend_apis.items(): + print(f"\n Backend: {backend_name}") + + # Capture inner tactics for this backend + with AutoTuner.get().capture() as inner_capture, torch.inference_mode(): + output = backend_api( + x_fp4, # input/act_fp4 + w_fp4, # weight + x_sf_block, # input_scale/act_sf + w_sf_block, # weight_scale + alpha_tensor, # alpha + dtype # output_dtype + ) - # Test CuteDSL backend with autotuning - with torch.inference_mode(), autotune(): - output_cutedsl = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cutedsl') + inner_tactics_list = list(inner_capture) + print(f" Found {len(inner_tactics_list)} inner tactics") + + # Verify tactics uniqueness (ensure we're testing different tactics, not repeating the same one) + tactic_values = [t[0][1] for t in inner_tactics_list] + unique_tactics = len(set(tactic_values)) + assert len(tactic_values) == unique_tactics, \ + f"Duplicate tactics detected! Total: {len(tactic_values)}, Unique: {unique_tactics}" + + # Test each tactic immediately (while _last_capture is still valid) + for tactic_idx, inner_tactic in enumerate(inner_tactics_list): + inner_runner, inner_tactic_value = inner_tactic[0] + runner_name = inner_runner.__class__.__name__ + + # Replay this tactic + with AutoTuner.get().replay(inner_tactic), torch.inference_mode(): + # Call backend API directly (using positional args) + output = backend_api( + x_fp4, # input/act_fp4 + w_fp4, # weight + x_sf_block, # input_scale/act_sf + w_sf_block, # weight_scale + alpha_tensor, # alpha + dtype # output_dtype + ) - # Verify CuteDSL result matches reference - torch.cuda.synchronize() - torch.testing.assert_close(output_cutedsl, output_ref, rtol=1e-2, atol=0.15) + # Verify correctness + torch.testing.assert_close(output, + output_ref, + rtol=1e-2, + atol=0.15) - # Capture all tactics using AutoTuner.capture() - with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): - output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cutedsl') - - # Convert to list and filter CuteDSL tactics - all_tactics_list = list(all_tactics) - cutedsl_tactics = [ - t for t in all_tactics_list if 'CuteDSL' in t[0][0].__class__.__name__ - ] + total_tactics_tested += 1 + tactics_by_backend[runner_name].append(total_tactics_tested) + print(f" ✓ Tactic {tactic_idx+1}/{len(inner_tactics_list)}: " + f"{runner_name} tactic={inner_tactic_value} - PASSED") - print(f"\n{'='*80}") - print( - f"Testing CuteDSL tactics for M={SEQ_LEN}, N={OUTPUT_SIZE}, K={HIDDEN_SIZE}" - ) - print(f"Total CuteDSL tactics: {len(cutedsl_tactics)}") - print(f"{'='*80}") + # Step 4: Test cuda_core if it's available (single tactic, no capture needed) + if test_cuda_core: + print(f"\n Backend: cuda_core") + print(f" Found 1 tactic (single implementation, no autotuning)") - # Test each CuteDSL tactic - for idx, tactic in enumerate(cutedsl_tactics): - with AutoTuner.get().replay(tactic), torch.inference_mode(): - output = torch.ops.trtllm.nvfp4_gemm_unified( + with torch.inference_mode(): + output_cuda_core = torch.ops.trtllm.nvfp4_gemm_unified( act_fp4=x_fp4, weight=w_fp4, act_sf=x_sf_block, @@ -654,16 +606,28 @@ def test_nvfp4_gemm_unified_cutedsl_all_tactics(dtype, mnk): alpha=alpha_tensor, output_dtype=dtype, to_userbuffers=False, - backend='cutedsl') + backend='cuda_core') - torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) - runner, tactic_value = tactic[0] - print( - f" ✓ CuteDSL tactic {idx+1}/{len(cutedsl_tactics)}: {runner.__class__.__name__} tactic={tactic_value} - PASSED" - ) + torch.testing.assert_close(output_cuda_core, + output_ref, + rtol=1e-2, + atol=0.15) - print(f"{'='*80}") - print(f"All {len(cutedsl_tactics)} CuteDSL tactics verified successfully!") + total_tactics_tested += 1 + tactics_by_backend['CudaCoreNVFP4Runner'].append(total_tactics_tested) + print(f" ✓ Tactic 1/1: CudaCoreNVFP4Runner tactic=0 - PASSED") + + print(f"\n{'='*80}") + print(f"All {total_tactics_tested} tactics verified successfully!") + print(f"\nBreakdown by backend:") + for runner_name, indices in tactics_by_backend.items(): + print(f" - {runner_name}: {len(indices)} tactics") + if test_cuda_core: + print(f"\n Note: cuda_core has no autotuning (single tactic)") + print(f" Note: Tested all inner layer tactics for each backend") + print( + f" Outer layer (backend selection) was tested separately with backend='auto'" + ) print(f"{'='*80}\n") From c6c2eea4a51f0a24c87d4ddb68ae14268429d602 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Fri, 21 Nov 2025 09:16:34 +0000 Subject: [PATCH 11/21] rename nvfp4_gemm to nvfp4_gemm_cutlass and rename nvfp4_gemm_unified to nvfp4_gemm Signed-off-by: Shijie Wang --- .../torch_compile_and_piecewise_cuda_graph.md | 2 +- .../_torch/auto_deploy/custom_ops/quant.py | 2 +- .../multi_stream/auto_multi_stream.py | 2 +- .../compilation/patterns/ar_residual_norm.py | 12 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 4 +- .../_torch/custom_ops/torch_custom_ops.py | 38 +++--- .../modules/fused_moe/fused_moe_cute_dsl.py | 7 +- tensorrt_llm/_torch/modules/linear.py | 8 +- .../_torch/thop/parallel/test_fp4_linear.py | 109 +++++++++--------- 9 files changed, 89 insertions(+), 95 deletions(-) diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 786ab39b51fd..97bc8a96290b 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -221,7 +221,7 @@ To address this, we implemented an auto multi-stream scheduler: call_function fp4_quantize_2 trtllm.fp4_quantize.default (mm_1, arg18_1, 16) {} call_function getitem_9 (fp4_quantize_2, 0) {} call_function getitem_10 (fp4_quantize_2, 1) {} - call_function nvfp4_gemm_2 trtllm.nvfp4_gemm.default (getitem_9, arg19_1, getitem_10, arg20_1, arg21_1, torch.bfloat16) {} + call_function nvfp4_gemm_2 trtllm.nvfp4_gemm_cutlass.default (getitem_9, arg19_1, getitem_10, arg20_1, arg21_1, torch.bfloat16) {} call_function permute_2 aten.permute.default (arg17_1, [1, 0]) {} call_function record_event_1 trtllm.record_event (0,) {} call_function silu_and_mul_1 trtllm.silu_and_mul.default (nvfp4_gemm_2,) {} diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py index d892cf6417ba..22e0376e1d65 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py @@ -337,7 +337,7 @@ def nvfp4_linear( input, input_scale, TRTLLM_NVFP4_SCALING_VECTOR_SIZE, False ) with autotune(): - output = torch.ops.trtllm.nvfp4_gemm( + output = torch.ops.trtllm.nvfp4_gemm_cutlass( x_fp4, weight_fp4, x_sf_block, weight_scale, alpha, input.dtype ) diff --git a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py index c2d3cf012a05..73693e3f2af5 100644 --- a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py +++ b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py @@ -43,7 +43,7 @@ def estimate_time(node: Node) -> int: gemm_ops = { torch.ops.aten.mm.default, - torch.ops.trtllm.nvfp4_gemm.default, + torch.ops.trtllm.nvfp4_gemm_cutlass.default, torch.ops.trtllm.fp8_batched_gemm_trtllmgen.default, torch.ops.trtllm.w4a8_mxfp4_fp8_gemm.default, torch.ops.trtllm.finegrained_mixed_dtype_gemm.default, diff --git a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py index afbaa0949df3..fa3632ea642f 100644 --- a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py +++ b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py @@ -555,10 +555,10 @@ def target_scaled_mm_prologue_pattern( def register_nvfp4_gemm_prologue(custom_pass: PatternMatcherPass): trtllm_nvfp4_gemm_default = CallFunction( - torch.ops.trtllm.nvfp4_gemm.default, KeywordArg('act_fp4'), - KeywordArg('weight'), KeywordArg('act_sf'), - KeywordArg('weight_scale'), KeywordArg('alpha'), - KeywordArg('output_dtype')) + torch.ops.trtllm.nvfp4_gemm_cutlass.default, + KeywordArg('act_fp4'), KeywordArg('weight'), + KeywordArg('act_sf'), KeywordArg('weight_scale'), + KeywordArg('alpha'), KeywordArg('output_dtype')) ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, trtllm_nvfp4_gemm_default) @@ -580,12 +580,12 @@ def target_nvfp4_gemm_prologue_pattern( alpha: torch.Tensor, output_dtype: torch.dtype, ): - nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm( + nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm_cutlass( act_fp4, weight, act_sf, weight_scale, alpha, output_dtype, True) return nvfp4_gemm_output - # No extra check needed as the output dtype of nvfp4_gemm has been verified when + # No extra check needed as the output dtype of nvfp4_gemm_cutlass has been verified when # ub_copy is inserted. register_replacement( empty_nvfp4_gemm_prologue_pattern, 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 8355bb58c3c4..d74fd6304b82 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -365,8 +365,8 @@ def cute_dsl_nvfp4_gemm_blackwell( """CuteDSL-based NVFP4 GEMM optimized for Blackwell. Note: - This function is primarily used internally by nvfp4_gemm_unified. - Direct usage is discouraged. Consider using nvfp4_gemm_unified instead + 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. """ tuner = AutoTuner.get() diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 9245d7b5166e..135e0c4fc9b1 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -602,8 +602,8 @@ def nvfp4_gemm_cublaslt( """cuBLASLt-based NVFP4 GEMM with heuristic-based auto-tuning. Note: - This function is primarily used internally by nvfp4_gemm_unified. - Direct usage is discouraged. Consider using nvfp4_gemm_unified instead + 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. """ tuner = AutoTuner.get() @@ -642,8 +642,8 @@ def _( dtype=output_dtype) -@torch.library.custom_op("trtllm::nvfp4_gemm", mutates_args=()) -def nvfp4_gemm( +@torch.library.custom_op("trtllm::nvfp4_gemm_cutlass", mutates_args=()) +def nvfp4_gemm_cutlass( act_fp4: torch.Tensor, weight: torch.Tensor, act_sf: torch.Tensor, @@ -655,8 +655,8 @@ def nvfp4_gemm( """CUTLASS-based NVFP4 GEMM with auto-tuning. Note: - This function is primarily used internally by nvfp4_gemm_unified. - Direct usage is discouraged. Consider using nvfp4_gemm_unified instead + 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. """ tuner = AutoTuner.get() @@ -678,7 +678,7 @@ def nvfp4_gemm( tactic=best_tactic) -@nvfp4_gemm.register_fake +@nvfp4_gemm_cutlass.register_fake def _( act_fp4: torch.Tensor, weight: torch.Tensor, @@ -833,10 +833,10 @@ def forward( out_dtype=self.output_dtype, to_userbuffers=self.to_userbuffers) elif tactic == "cutlass": - return torch.ops.trtllm.nvfp4_gemm(act_fp4, weight, act_sf, - weight_scale, alpha, - self.output_dtype, - self.to_userbuffers) + return torch.ops.trtllm.nvfp4_gemm_cutlass(act_fp4, weight, act_sf, + weight_scale, alpha, + self.output_dtype, + self.to_userbuffers) elif tactic == "cublaslt": return torch.ops.trtllm.nvfp4_gemm_cublaslt(act_fp4, weight, act_sf, weight_scale, alpha, @@ -846,16 +846,16 @@ def forward( return torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( act_fp4, weight, act_sf, weight_scale, alpha, self.output_dtype) elif tactic == -1: - return torch.ops.trtllm.nvfp4_gemm(act_fp4, weight, act_sf, - weight_scale, alpha, - self.output_dtype, - self.to_userbuffers) + return torch.ops.trtllm.nvfp4_gemm_cutlass(act_fp4, weight, act_sf, + weight_scale, alpha, + self.output_dtype, + self.to_userbuffers) else: raise ValueError(f"Invalid tactic: {tactic}") -@torch.library.custom_op("trtllm::nvfp4_gemm_unified", mutates_args=()) -def nvfp4_gemm_unified( +@torch.library.custom_op("trtllm::nvfp4_gemm", mutates_args=()) +def nvfp4_gemm( act_fp4: torch.Tensor, weight: torch.Tensor, act_sf: torch.Tensor, @@ -916,7 +916,7 @@ def nvfp4_gemm_unified( try: _, best_tactic = tuner.choose_one( - "trtllm::nvfp4_gemm_unified::gemm", + "trtllm::nvfp4_gemm::gemm", [runner], FP4GemmRunner. tuning_config, # All runners use the same tuning_config @@ -940,7 +940,7 @@ def nvfp4_gemm_unified( ) -@nvfp4_gemm_unified.register_fake +@nvfp4_gemm.register_fake def _( act_fp4: torch.Tensor, weight: torch.Tensor, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py index 5eb32965d815..4e355734d375 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py @@ -143,9 +143,10 @@ def cute_dsl_nvfp4_grouped_gemm_ref( a_sliced = a[start * tile_size:end * tile_size] a_sf_sliced = a_sf[start * tile_size * k // scaling_vector_size:end * tile_size * k // scaling_vector_size] - ref[start * tile_size:end * tile_size] = torch.ops.trtllm.nvfp4_gemm( - a_sliced.view(torch.uint8), b[i].view(torch.uint8), a_sf_sliced, - b_sf[i], alpha[i], output_dtype) + ref[start * tile_size:end * + tile_size] = torch.ops.trtllm.nvfp4_gemm_cutlass( + a_sliced.view(torch.uint8), b[i].view(torch.uint8), a_sf_sliced, + b_sf[i], alpha[i], output_dtype) return ref diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 1d7bd99e4d23..a050b2819d12 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -918,11 +918,9 @@ def apply(self, module: Linear, input: torch.Tensor, backend = getattr(module, 'nvfp4_backend', 'auto') # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL - output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4, module.weight, - act_sf, - module.weight_scale, - module.alpha, module.dtype, - False, backend) + output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, act_sf, + module.weight_scale, module.alpha, + module.dtype, False, backend) # Take the dim of out_features if padded. Make sure the output is contiguous if output.shape[-1] > module.out_features: output = output[..., :module.out_features].contiguous() diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py index 7f1a4ccb581f..a549b52fa4ec 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_linear.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_linear.py @@ -325,7 +325,7 @@ def nvfp4_gemm_perf_test( f"ref tune, m={SEQ_LEN}, k={HIDDEN_SIZE}, n={OUTPUT_SIZE}", color="orange"): with torch.inference_mode(), autotune(): - output_ref = torch.ops.trtllm.nvfp4_gemm( + output_ref = torch.ops.trtllm.nvfp4_gemm_cutlass( x_fp4, w_fp4, x_sf_block, w_sf_block, alpha_tensor, dtype) torch.testing.assert_close(output, output_ref) print(f"PASSED") @@ -368,7 +368,7 @@ def nvfp4_gemm_perf_test( f"ref warmup, m={SEQ_LEN}, k={HIDDEN_SIZE}, n={OUTPUT_SIZE}", color="red"): for _ in range(warmup_iterations): - output_ref = torch.ops.trtllm.nvfp4_gemm( + output_ref = torch.ops.trtllm.nvfp4_gemm_cutlass( x_fp4_list[buffer_idx % workspace_count], w_fp4_list[buffer_idx % workspace_count], x_sf_block_list[buffer_idx % workspace_count], @@ -381,7 +381,7 @@ def nvfp4_gemm_perf_test( f"ref run, m={SEQ_LEN}, k={HIDDEN_SIZE}, n={OUTPUT_SIZE}", color="red"): for i in range(iterations): - output_ref = torch.ops.trtllm.nvfp4_gemm( + output_ref = torch.ops.trtllm.nvfp4_gemm_cutlass( x_fp4_list[buffer_idx % workspace_count], w_fp4_list[buffer_idx % workspace_count], x_sf_block_list[buffer_idx % workspace_count], @@ -429,7 +429,7 @@ def nvfp4_gemm_perf_test( (128, 16384, 128, "Wide N: N >> K"), ]) def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): - """Test nvfp4_gemm_unified with auto backend selection, ensuring all tactics are tested.""" + """Test nvfp4_gemm with auto backend selection, ensuring all tactics are tested.""" from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.cublaslt_utils import IS_CUBLASLT_AVAILABLE @@ -459,27 +459,25 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): # Reference: Use CUTLASS backend explicitly for reference output with torch.inference_mode(): - output_ref = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cutlass') + output_ref = torch.ops.trtllm.nvfp4_gemm(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutlass') # Test auto backend selection with autotuning with torch.inference_mode(), autotune(): - output_auto = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='auto') + output_auto = torch.ops.trtllm.nvfp4_gemm(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='auto') AutoTuner.get().print_profiling_cache() @@ -488,25 +486,25 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): torch.testing.assert_close(output_auto, output_ref, rtol=1e-2, atol=0.15) # Test all combinations of outer layer (backend selection) and inner layer (backend tactics) - # Outer layer: nvfp4_gemm_unified selects backend + # Outer layer: nvfp4_gemm selects backend # Inner layer: each backend has its own tactics from collections import defaultdict print(f"\n{'='*80}") - print(f"Testing nvfp4_gemm_unified (2-layer tactics): {desc}") + print(f"Testing nvfp4_gemm (2-layer tactics): {desc}") print(f"Shape: M={SEQ_LEN}, K={HIDDEN_SIZE}, N={OUTPUT_SIZE}") print(f"{'='*80}") print(f"\n[Outer Layer] Capturing backend selection tactics...") with AutoTuner.get().capture() as outer_capture, torch.inference_mode(): - output = torch.ops.trtllm.nvfp4_gemm_unified(act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='auto') + output = torch.ops.trtllm.nvfp4_gemm(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='auto') outer_tactics_list = list(outer_capture) print(f" Found {len(outer_tactics_list)} outer layer tactics (backends)") @@ -524,7 +522,7 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): backend_apis = {} if IS_CUTLASS_DSL_AVAILABLE: if 'cutlass' in backend_map: - backend_apis['cutlass'] = torch.ops.trtllm.nvfp4_gemm + backend_apis['cutlass'] = torch.ops.trtllm.nvfp4_gemm_cutlass if IS_CUBLASLT_AVAILABLE: if 'cublaslt' in backend_map: backend_apis['cublaslt'] = torch.ops.trtllm.nvfp4_gemm_cublaslt @@ -598,7 +596,7 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): print(f" Found 1 tactic (single implementation, no autotuning)") with torch.inference_mode(): - output_cuda_core = torch.ops.trtllm.nvfp4_gemm_unified( + output_cuda_core = torch.ops.trtllm.nvfp4_gemm( act_fp4=x_fp4, weight=w_fp4, act_sf=x_sf_block, @@ -640,7 +638,7 @@ def test_nvfp4_gemm_unified_all_tactics(dtype, mnk): (128, 2112, 7168), (128, 4096, 7168), (128, 7168, 2048), [127, 1024, 3200]]) def test_fp4_linear_cublaslt(dtype, mnk): - """Test cuBLASLt FP4 GEMM implementation and compare with nvfp4_gemm""" + """Test cuBLASLt FP4 GEMM implementation and compare with nvfp4_gemm_cutlass""" from tensorrt_llm._torch.cublaslt_utils import IS_CUBLASLT_AVAILABLE if not IS_CUBLASLT_AVAILABLE: pytest.skip("cuBLASLt FP4 GEMM not available in this build") @@ -674,11 +672,10 @@ def test_fp4_linear_cublaslt(dtype, mnk): alpha=alpha_tensor, output_dtype=dtype) - # Reference implementation: use torch.ops.trtllm.nvfp4_gemm (CUTLASS) + # Reference implementation: use torch.ops.trtllm.nvfp4_gemm_cutlass (CUTLASS) with torch.inference_mode(): - output_cutlass = torch.ops.trtllm.nvfp4_gemm(x_fp4, w_fp4, x_sf_block, - w_sf_block, alpha_ref, - dtype) + output_cutlass = torch.ops.trtllm.nvfp4_gemm_cutlass( + x_fp4, w_fp4, x_sf_block, w_sf_block, alpha_ref, dtype) # Compare results torch.cuda.synchronize() @@ -715,26 +712,24 @@ def test_fp4_linear_cuda_core(dtype, mnk): alpha_tensor = torch.tensor(alpha_ref, dtype=torch.float32).cuda() # Reference: Use CUTLASS backend - output_ref = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cutlass') + output_ref = torch.ops.trtllm.nvfp4_gemm(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cutlass') # Test CUDA Core backend - output_cuda_core = torch.ops.trtllm.nvfp4_gemm_unified( - act_fp4=x_fp4, - weight=w_fp4, - act_sf=x_sf_block, - weight_scale=w_sf_block, - alpha=alpha_tensor, - output_dtype=dtype, - to_userbuffers=False, - backend='cuda_core') + output_cuda_core = torch.ops.trtllm.nvfp4_gemm(act_fp4=x_fp4, + weight=w_fp4, + act_sf=x_sf_block, + weight_scale=w_sf_block, + alpha=alpha_tensor, + output_dtype=dtype, + to_userbuffers=False, + backend='cuda_core') # Compare results torch.cuda.synchronize() From 1fcbd6cf7c68478c98f5cd5f3cbe8758bcc5cc3e Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Tue, 25 Nov 2025 07:21:29 +0000 Subject: [PATCH 12/21] minor change Signed-off-by: Shijie Wang --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 25 ++++++-- .../_torch/custom_ops/torch_custom_ops.py | 59 +++++++++++-------- 2 files changed, 56 insertions(+), 28 deletions(-) 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 d74fd6304b82..cad17098e595 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -34,7 +34,7 @@ Sm100BlockScaledPersistentDenseGemmKernel from ..cute_dsl_kernels.blackwell.utils import make_ptr - class CuteDSLNVFP4BlackwellRunner(TunableRunner): + class CuteDSLNVFP4BlackwellLinear(TunableRunner): kernel_class = Sm100BlockScaledPersistentDenseGemmKernel kernel_cache = dict() tuning_config = TuningConfig( @@ -48,10 +48,7 @@ class CuteDSLNVFP4BlackwellRunner(TunableRunner): def __init__(self, output_dtype: torch.dtype): super().__init__() - if get_sm_version() not in [100, 103]: - raise ValueError( - f"SM version {get_sm_version()} is not supported for {self.__class__.__name__}, it only supports SM 100 and SM 103" - ) + if output_dtype != torch.bfloat16: raise ValueError( @@ -69,6 +66,15 @@ def get_valid_tactics( profile: OptimizationProfile, **kwargs, ) -> List[Tuple[int, int]]: + # Early exit: Check SM version - CuteDSL NVFP4 only supports SM 100 and SM 103 + sm_version = get_sm_version() + if sm_version not in [100, 103]: + logger.debug( + f"CuteDSL: SM version {sm_version} is not supported. " + f"CuteDSL NVFP4 only supports SM 100 (B200) and SM 103 (B300). Skipping all tactics." + ) + return [] + assert inputs[0].dim() == 2 assert inputs[1].dim() == 2 @@ -369,6 +375,14 @@ def cute_dsl_nvfp4_gemm_blackwell( Direct usage is discouraged. Consider using nvfp4_gemm instead for automatic backend selection with better performance. """ + # Validate SM version before attempting to use CuteDSL + sm_version = get_sm_version() + if sm_version not in [100, 103]: + raise ValueError( + f"CuteDSL NVFP4 backend requires SM 100 (B200) or SM 103 (B300), but got SM {sm_version}. " + f"Please use nvfp4_gemm with backend='auto' for automatic backend selection." + ) + tuner = AutoTuner.get() runner = CuteDSLNVFP4BlackwellLinear(output_dtype) @@ -378,6 +392,7 @@ def cute_dsl_nvfp4_gemm_blackwell( [runner], runner.__class__.tuning_config, inputs, + ) output = runner(inputs, tactic=best_tactic) return output diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 135e0c4fc9b1..3ff27ff92f05 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -748,29 +748,42 @@ def get_valid_tactics(self, # Add CuteDSL runner if available if backend in ["auto", "cutedsl"]: if IS_CUTLASS_DSL_AVAILABLE: - # Check if CuteDSL actually supports the current shape - from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ - CuteDSLNVFP4BlackwellLinear - cutedsl_runner = CuteDSLNVFP4BlackwellLinear(self.output_dtype) - cutedsl_tactics = cutedsl_runner.get_valid_tactics( - inputs, profile) - - if cutedsl_tactics: - # CuteDSL supports this shape - tactics.append("cutedsl") - elif backend == "cutedsl": - # Explicitly requested CuteDSL but it doesn't support this shape - m, n, k = inputs[0].shape[0], inputs[1].shape[ - 0], inputs[0].shape[1] * 2 - raise ValueError( - f"CuteDSL backend does not support the current shape:\n" - f" M={m}, N={n}, K={k}\n" - f"CuteDSL requires 16-byte alignment for major (contiguous) dimensions:\n" - f" - K must be divisible by 32 (FP4 K-major layout): K%32={'0✓' if k % 32 == 0 else str(k%32)+'✗'}\n" - f" - Or the combination of (M, N, K, tiling, cluster shape) is not supported\n" - f"Please use backend='auto' to automatically select a compatible backend." - ) - # else: backend='auto' and CuteDSL doesn't support → silently skip + # Check SM version first - CuteDSL NVFP4 only supports SM 100 (B200) + sm_version = get_sm_version() + if sm_version != 100: + if backend == "cutedsl": + # Explicitly requested CuteDSL but SM version not supported + raise ValueError( + f"CuteDSL NVFP4 backend requires SM 100 (B200), but got SM {sm_version}. " + f"CuteDSL NVFP4 is not supported on this GPU architecture. " + f"Please use backend='auto' to automatically select a compatible backend." + ) + # else: backend='auto' → silently skip CuteDSL + else: + # SM version OK, check if CuteDSL supports the current shape + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \ + CuteDSLNVFP4BlackwellLinear + cutedsl_runner = CuteDSLNVFP4BlackwellLinear( + self.output_dtype) + cutedsl_tactics = cutedsl_runner.get_valid_tactics( + inputs, profile) + + if cutedsl_tactics: + # CuteDSL supports this shape + tactics.append("cutedsl") + elif backend == "cutedsl": + # Explicitly requested CuteDSL but it doesn't support this shape + m, n, k = inputs[0].shape[0], inputs[1].shape[ + 0], inputs[0].shape[1] * 2 + raise ValueError( + f"CuteDSL backend does not support the current shape:\n" + f" M={m}, N={n}, K={k}\n" + f"CuteDSL requires 16-byte alignment for major (contiguous) dimensions:\n" + f" - K must be divisible by 32 (FP4 K-major layout): K%32={'0✓' if k % 32 == 0 else str(k%32)+'✗'}\n" + f" - Or the combination of (M, N, K, tiling, cluster shape) is not supported\n" + f"Please use backend='auto' to automatically select a compatible backend." + ) + # else: backend='auto' and CuteDSL doesn't support shape → silently skip elif backend == "cutedsl": raise ValueError( "CuteDSL backend is not available. " From a1cfc8244c36a894b7b8765b523540c21b5d5354 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Wed, 26 Nov 2025 10:23:20 +0000 Subject: [PATCH 13/21] minor changes Signed-off-by: Shijie Wang --- .../features/torch_compile_and_piecewise_cuda_graph.md | 2 +- .../_torch/compilation/multi_stream/auto_multi_stream.py | 2 +- .../_torch/compilation/patterns/ar_residual_norm.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 97bc8a96290b..786ab39b51fd 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -221,7 +221,7 @@ To address this, we implemented an auto multi-stream scheduler: call_function fp4_quantize_2 trtllm.fp4_quantize.default (mm_1, arg18_1, 16) {} call_function getitem_9 (fp4_quantize_2, 0) {} call_function getitem_10 (fp4_quantize_2, 1) {} - call_function nvfp4_gemm_2 trtllm.nvfp4_gemm_cutlass.default (getitem_9, arg19_1, getitem_10, arg20_1, arg21_1, torch.bfloat16) {} + call_function nvfp4_gemm_2 trtllm.nvfp4_gemm.default (getitem_9, arg19_1, getitem_10, arg20_1, arg21_1, torch.bfloat16) {} call_function permute_2 aten.permute.default (arg17_1, [1, 0]) {} call_function record_event_1 trtllm.record_event (0,) {} call_function silu_and_mul_1 trtllm.silu_and_mul.default (nvfp4_gemm_2,) {} diff --git a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py index 73693e3f2af5..c2d3cf012a05 100644 --- a/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py +++ b/tensorrt_llm/_torch/compilation/multi_stream/auto_multi_stream.py @@ -43,7 +43,7 @@ def estimate_time(node: Node) -> int: gemm_ops = { torch.ops.aten.mm.default, - torch.ops.trtllm.nvfp4_gemm_cutlass.default, + torch.ops.trtllm.nvfp4_gemm.default, torch.ops.trtllm.fp8_batched_gemm_trtllmgen.default, torch.ops.trtllm.w4a8_mxfp4_fp8_gemm.default, torch.ops.trtllm.finegrained_mixed_dtype_gemm.default, diff --git a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py index fa3632ea642f..e4b20903b2e6 100644 --- a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py +++ b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py @@ -555,10 +555,10 @@ def target_scaled_mm_prologue_pattern( def register_nvfp4_gemm_prologue(custom_pass: PatternMatcherPass): trtllm_nvfp4_gemm_default = CallFunction( - torch.ops.trtllm.nvfp4_gemm_cutlass.default, - KeywordArg('act_fp4'), KeywordArg('weight'), - KeywordArg('act_sf'), KeywordArg('weight_scale'), - KeywordArg('alpha'), KeywordArg('output_dtype')) + torch.ops.trtllm.nvfp4_gemm.default, KeywordArg('act_fp4'), + KeywordArg('weight'), KeywordArg('act_sf'), + KeywordArg('weight_scale'), KeywordArg('alpha'), + KeywordArg('output_dtype')) ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, trtllm_nvfp4_gemm_default) From 22d31308ed9b8cfe9974132eb4bbb60a6ab00355 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Wed, 26 Nov 2025 11:58:37 +0000 Subject: [PATCH 14/21] support to_userbuffers in cutedsl nvfp4 gemm Signed-off-by: Shijie Wang --- .../_torch/auto_deploy/custom_ops/quant.py | 2 +- .../compilation/patterns/ar_residual_norm.py | 4 +- .../_torch/custom_ops/cute_dsl_custom_ops.py | 41 ++++++++++++++++--- .../_torch/custom_ops/torch_custom_ops.py | 3 +- .../modules/fused_moe/fused_moe_cute_dsl.py | 7 ++-- 5 files changed, 43 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py index 22e0376e1d65..d892cf6417ba 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py @@ -337,7 +337,7 @@ def nvfp4_linear( input, input_scale, TRTLLM_NVFP4_SCALING_VECTOR_SIZE, False ) with autotune(): - output = torch.ops.trtllm.nvfp4_gemm_cutlass( + output = torch.ops.trtllm.nvfp4_gemm( x_fp4, weight_fp4, x_sf_block, weight_scale, alpha, input.dtype ) diff --git a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py index e4b20903b2e6..afbaa0949df3 100644 --- a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py +++ b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py @@ -580,12 +580,12 @@ def target_nvfp4_gemm_prologue_pattern( alpha: torch.Tensor, output_dtype: torch.dtype, ): - nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm_cutlass( + nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm( act_fp4, weight, act_sf, weight_scale, alpha, output_dtype, True) return nvfp4_gemm_output - # No extra check needed as the output dtype of nvfp4_gemm_cutlass has been verified when + # No extra check needed as the output dtype of nvfp4_gemm has been verified when # ub_copy is inserted. register_replacement( empty_nvfp4_gemm_prologue_pattern, 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 cad17098e595..5654810465bd 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -45,7 +45,9 @@ class CuteDSLNVFP4BlackwellLinear(TunableRunner): use_cold_l2_cache=True, ) - def __init__(self, output_dtype: torch.dtype): + def __init__(self, + output_dtype: torch.dtype, + to_userbuffers: bool = False): super().__init__() @@ -55,10 +57,19 @@ def __init__(self, output_dtype: torch.dtype): f"CuteDSL NVFP4 only supports bfloat16 output, got {output_dtype}" ) self.output_dtype = output_dtype + self.to_userbuffers = to_userbuffers # rewrite the hash function because the value of self.alpha doesn't affect the tactic. def unique_id(self): - return (self.output_dtype, ) + return (self.output_dtype, self.to_userbuffers) + + def __hash__(self): + return hash((self.output_dtype, self.to_userbuffers)) + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return False + return self.output_dtype == other.output_dtype and self.to_userbuffers == other.to_userbuffers def get_valid_tactics( self, @@ -228,9 +239,16 @@ def forward( a_tensor, b_tensor, a_sf_tensor, b_sf_tensor, alpha_tensor = inputs m, k, n = a_tensor.shape[0], a_tensor.shape[1], b_tensor.shape[0] - c_tensor = torch.empty(*(m, n), - dtype=self.output_dtype, - device="cuda") + + # Allocate output tensor from UserBuffers or regular CUDA memory + if self.to_userbuffers: + from tensorrt_llm.bindings import torch_ext + c_tensor, _ = torch_ext.create_userbuffers_tensor( + [m, n], self.output_dtype) + else: + c_tensor = torch.empty(*(m, n), + dtype=self.output_dtype, + device="cuda") if swap_ab: c_tensor = c_tensor.permute(1, 0) @@ -367,9 +385,19 @@ def cute_dsl_nvfp4_gemm_blackwell( weight_scale: torch.Tensor, alpha: torch.Tensor, output_dtype: torch.dtype, + to_userbuffers: bool = False, ) -> torch.Tensor: """CuteDSL-based NVFP4 GEMM optimized for Blackwell. + 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_dtype: Output data type (must be bfloat16) + to_userbuffers: Whether to allocate output from UserBuffers pool + Note: This function is primarily used internally by nvfp4_gemm. Direct usage is discouraged. Consider using nvfp4_gemm instead @@ -385,7 +413,7 @@ def cute_dsl_nvfp4_gemm_blackwell( tuner = AutoTuner.get() - runner = CuteDSLNVFP4BlackwellLinear(output_dtype) + runner = CuteDSLNVFP4BlackwellLinear(output_dtype, to_userbuffers) inputs = [input, weight, input_scale, weight_scale, alpha] _, best_tactic = tuner.choose_one( "trtllm::cute_dsl_nvfp4_gemm_blackwell", @@ -405,6 +433,7 @@ def _( weight_scale: torch.Tensor, alpha: torch.Tensor, # Match custom op signature output_dtype: torch.dtype, + to_userbuffers: bool = False, ): # [m, k] shape = list(mat_a.shape) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 3ff27ff92f05..1ff54a0ff24c 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -857,7 +857,8 @@ def forward( self.to_userbuffers) elif tactic == "cutedsl": return torch.ops.trtllm.cute_dsl_nvfp4_gemm_blackwell( - act_fp4, weight, act_sf, weight_scale, alpha, self.output_dtype) + act_fp4, weight, act_sf, weight_scale, alpha, self.output_dtype, + self.to_userbuffers) elif tactic == -1: return torch.ops.trtllm.nvfp4_gemm_cutlass(act_fp4, weight, act_sf, weight_scale, alpha, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py index 4e355734d375..5eb32965d815 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl.py @@ -143,10 +143,9 @@ def cute_dsl_nvfp4_grouped_gemm_ref( a_sliced = a[start * tile_size:end * tile_size] a_sf_sliced = a_sf[start * tile_size * k // scaling_vector_size:end * tile_size * k // scaling_vector_size] - ref[start * tile_size:end * - tile_size] = torch.ops.trtllm.nvfp4_gemm_cutlass( - a_sliced.view(torch.uint8), b[i].view(torch.uint8), a_sf_sliced, - b_sf[i], alpha[i], output_dtype) + ref[start * tile_size:end * tile_size] = torch.ops.trtllm.nvfp4_gemm( + a_sliced.view(torch.uint8), b[i].view(torch.uint8), a_sf_sliced, + b_sf[i], alpha[i], output_dtype) return ref From 83646782cc217a1f3a48a0914c7e89ebd06f0a66 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Fri, 28 Nov 2025 05:35:35 +0000 Subject: [PATCH 15/21] Support to_userbuffers in CuteDSL and fix nvfp4_gemm pattern matching Signed-off-by: Shijie Wang --- .../compilation/patterns/ar_residual_norm.py | 58 ++++++++++++++++--- .../_torch/custom_ops/cute_dsl_custom_ops.py | 3 +- tensorrt_llm/_torch/modules/linear.py | 11 +++- 3 files changed, 59 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py index afbaa0949df3..34919d86bfdc 100644 --- a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py +++ b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py @@ -554,11 +554,24 @@ def target_scaled_mm_prologue_pattern( ) def register_nvfp4_gemm_prologue(custom_pass: PatternMatcherPass): + act_fp4_key = KeywordArg('act_fp4') + weight_key = KeywordArg('weight') + act_sf_key = KeywordArg('act_sf') + weight_scale_key = KeywordArg('weight_scale') + alpha_key = KeywordArg('alpha') + output_dtype_key = KeywordArg('output_dtype') + to_userbuffers_key = KeywordArg('to_userbuffers') + backend_key = KeywordArg('backend') trtllm_nvfp4_gemm_default = CallFunction( - torch.ops.trtllm.nvfp4_gemm.default, KeywordArg('act_fp4'), - KeywordArg('weight'), KeywordArg('act_sf'), - KeywordArg('weight_scale'), KeywordArg('alpha'), - KeywordArg('output_dtype')) + torch.ops.trtllm.nvfp4_gemm.default, + act_fp4_key, + weight_key, + act_sf_key, + weight_scale_key, + alpha_key, + output_dtype_key, + to_userbuffers=to_userbuffers_key, + backend=backend_key) ub_copy = CallFunction(torch.ops.trtllm.copy_to_userbuffers, trtllm_nvfp4_gemm_default) @@ -569,6 +582,8 @@ def empty_nvfp4_gemm_prologue_pattern( weight_scale: torch.Tensor, alpha: torch.Tensor, output_dtype: torch.dtype, + to_userbuffers: bool, + backend: str, ): return @@ -579,18 +594,45 @@ def target_nvfp4_gemm_prologue_pattern( weight_scale: torch.Tensor, alpha: torch.Tensor, output_dtype: torch.dtype, + to_userbuffers: bool, + backend: str, ): nvfp4_gemm_output = torch.ops.trtllm.nvfp4_gemm( act_fp4, weight, act_sf, weight_scale, alpha, output_dtype, - True) + True, backend) return nvfp4_gemm_output - # No extra check needed as the output dtype of nvfp4_gemm has been verified when - # ub_copy is inserted. + def extra_check(match: Match) -> bool: + # Validate backend value + backend_node = match.kwargs.get('backend') + if backend_node is None: + # No backend specified, use default - OK + return True + + valid_backends = {'auto', 'cutlass', 'cublaslt', 'cutedsl'} + + # Case 1: backend is a Node with metadata + if hasattr(backend_node, 'meta') and 'val' in backend_node.meta: + backend_value = backend_node.meta['val'] + if isinstance(backend_value, str): + return backend_value in valid_backends + return False # Invalid type + + # Case 2: backend is a constant Node in the graph + if hasattr(backend_node, 'target'): + return backend_node.target in valid_backends + + # Case 3: backend is a Python literal in kwargs + if isinstance(backend_node, str): + return backend_node in valid_backends + + # Unknown format - reject to be safe + return False + register_replacement( empty_nvfp4_gemm_prologue_pattern, target_nvfp4_gemm_prologue_pattern, - [], + [extra_check], fwd_only, custom_pass, search_fn_pattern=ub_copy, 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 5654810465bd..b8866643148f 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -242,8 +242,7 @@ def forward( # Allocate output tensor from UserBuffers or regular CUDA memory if self.to_userbuffers: - from tensorrt_llm.bindings import torch_ext - c_tensor, _ = torch_ext.create_userbuffers_tensor( + c_tensor, _ = torch.ops.trtllm.create_userbuffers_tensor( [m, n], self.output_dtype) else: c_tensor = torch.empty(*(m, n), diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index a050b2819d12..c505881ab0e1 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -918,9 +918,14 @@ def apply(self, module: Linear, input: torch.Tensor, backend = getattr(module, 'nvfp4_backend', 'auto') # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL - output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, act_sf, - module.weight_scale, module.alpha, - module.dtype, False, backend) + output = torch.ops.trtllm.nvfp4_gemm(act_fp4, + module.weight, + act_sf, + module.weight_scale, + module.alpha, + module.dtype, + to_userbuffers=False, + backend=backend) # Take the dim of out_features if padded. Make sure the output is contiguous if output.shape[-1] > module.out_features: output = output[..., :module.out_features].contiguous() From e2b8e24abb0f4c84d78d8a46100cbe27352c1a9f Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Sun, 30 Nov 2025 10:21:57 +0000 Subject: [PATCH 16/21] minor change Signed-off-by: Shijie Wang --- .../_torch/custom_ops/torch_custom_ops.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 1ff54a0ff24c..1fb0d7f227cd 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -695,19 +695,26 @@ def _( class NVFP4GemmUnifiedRunner(TunableRunner): runner_dict = dict() - def __init__(self, to_userbuffers: bool, output_dtype: torch.dtype): + def __init__(self, + to_userbuffers: bool, + output_dtype: torch.dtype, + backend: str = "auto"): super().__init__() self.to_userbuffers = to_userbuffers self.output_dtype = output_dtype + self.backend = backend + + def unique_id(self): + """Include backend in cache key to avoid sharing cache across backends.""" + return (self.to_userbuffers, self.output_dtype, self.backend) - def get_valid_tactics(self, - inputs: List[torch.Tensor], + def get_valid_tactics(self, inputs: List[torch.Tensor], profile: OptimizationProfile, - backend: str = "auto", **kwargs) -> List[Tuple]: # return valid nvfp4 gemm implementations tactics = [] act_fp4, weight, act_sf, weight_scale, alpha = inputs + backend = self.backend if backend in ["auto", "cuda_core"]: is_cuda_core_supported = False @@ -800,8 +807,7 @@ def forward( ) -> torch.Tensor: act_fp4, weight, act_sf, weight_scale, alpha = inputs - # Check if a specific backend was requested - requested_backend = kwargs.get('backend', 'auto') + requested_backend = self.backend # If a specific backend was requested (not 'auto') and we're using fallback tactic # This can happen on cache miss, where AutoTuner uses tactic=-1 as default @@ -812,8 +818,7 @@ def forward( # Get valid tactics for the requested backend from tensorrt_llm._torch.autotuner import OptimizationProfile valid_tactics = self.get_valid_tactics(inputs, - OptimizationProfile(), - backend=requested_backend) + OptimizationProfile()) if not valid_tactics or requested_backend not in valid_tactics: # Requested backend doesn't support this shape @@ -921,7 +926,7 @@ def nvfp4_gemm( f"Invalid backend '{backend}'. Must be one of {valid_backends}") # Build list of runners based on backend parameter - runner = NVFP4GemmUnifiedRunner(to_userbuffers, output_dtype) + runner = NVFP4GemmUnifiedRunner(to_userbuffers, output_dtype, backend) # Use AutoTuner to select best runner and tactic # - For 'auto' mode: compare across all backends, find global optimum @@ -935,7 +940,6 @@ def nvfp4_gemm( FP4GemmRunner. tuning_config, # All runners use the same tuning_config [act_fp4, weight, act_sf, weight_scale, alpha], - backend=backend, ) except IndexError as e: # Provide more helpful error message @@ -950,7 +954,6 @@ def nvfp4_gemm( return runner( inputs=[act_fp4, weight, act_sf, weight_scale, alpha], tactic=best_tactic, - backend=backend, ) From 205d297bba85f134f896d00a4db8f450dd6dd031 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Mon, 1 Dec 2025 06:41:26 +0000 Subject: [PATCH 17/21] resolve review comments Signed-off-by: Shijie Wang --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 4 +--- tensorrt_llm/_torch/modules/linear.py | 11 +++++++---- 2 files changed, 8 insertions(+), 7 deletions(-) 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 b8866643148f..fd927a0258e0 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -50,8 +50,6 @@ def __init__(self, to_userbuffers: bool = False): super().__init__() - - if output_dtype != torch.bfloat16: raise ValueError( f"CuteDSL NVFP4 only supports bfloat16 output, got {output_dtype}" @@ -242,7 +240,7 @@ def forward( # Allocate output tensor from UserBuffers or regular CUDA memory if self.to_userbuffers: - c_tensor, _ = torch.ops.trtllm.create_userbuffers_tensor( + c_tensor = torch.ops.trtllm.create_userbuffers_tensor( [m, n], self.output_dtype) else: c_tensor = torch.empty(*(m, n), diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index c505881ab0e1..68304e897879 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -914,9 +914,6 @@ def apply(self, module: Linear, input: torch.Tensor, act_fp4, act_sf = torch.ops.trtllm.fp4_quantize( input, module.input_scale, module.scaling_vector_size, False) - # Backend selection: 'auto' (default) | 'cutlass' | 'cublaslt' | 'cutedsl' - backend = getattr(module, 'nvfp4_backend', 'auto') - # Use unified interface - supports CUTLASS, cuBLASLt, CuteDSL output = torch.ops.trtllm.nvfp4_gemm(act_fp4, module.weight, @@ -925,7 +922,7 @@ def apply(self, module: Linear, input: torch.Tensor, module.alpha, module.dtype, to_userbuffers=False, - backend=backend) + backend=module.nvfp4_backend) # Take the dim of out_features if padded. Make sure the output is contiguous if output.shape[-1] > module.out_features: output = output[..., :module.out_features].contiguous() @@ -2000,6 +1997,12 @@ def __init__( fused_weight_shard_indices_mapping: Optional[dict] = None, nvfp4_backend: str = "auto", ): + """ + Args: + nvfp4_backend: Backend selection for NVFP4 GEMM operations. + Supported values: "auto", "cutlass", "cublaslt", "cutedsl". + Default is "auto" which automatically selects the best backend. + """ from ..distributed import AllReduce super().__init__() From eba26425780345bc3d914c3b179f8d32636d065d Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Mon, 1 Dec 2025 06:55:00 +0000 Subject: [PATCH 18/21] add ENV TRTLLM_NVFP4_GEMM_BACKEND to override nvfp4_backend in Linear Signed-off-by: Shijie Wang --- tensorrt_llm/_torch/modules/linear.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 68304e897879..1a1769bfde61 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -2002,6 +2002,7 @@ def __init__( nvfp4_backend: Backend selection for NVFP4 GEMM operations. Supported values: "auto", "cutlass", "cublaslt", "cutedsl". Default is "auto" which automatically selects the best backend. + Can be overridden via TRTLLM_NVFP4_GEMM_BACKEND environment variable. """ from ..distributed import AllReduce @@ -2021,7 +2022,21 @@ def __init__( self.use_cute_dsl_blockscaling_mm = use_cute_dsl_blockscaling_mm self.disable_deep_gemm = disable_deep_gemm self.fused_weight_shard_indices_mapping = fused_weight_shard_indices_mapping - self.nvfp4_backend = nvfp4_backend + + # Support environment variable override for nvfp4_backend + nvfp4_backend_value = os.environ.get('TRTLLM_NVFP4_GEMM_BACKEND', + nvfp4_backend) + + # Validate backend selection + valid_backends = {'auto', 'cutlass', 'cublaslt', 'cutedsl'} + if nvfp4_backend_value not in valid_backends: + raise ValueError( + f"Invalid nvfp4_backend: '{nvfp4_backend_value}'. " + f"Supported values are: {', '.join(sorted(valid_backends))}. " + f"Set via constructor argument or TRTLLM_NVFP4_GEMM_BACKEND environment variable." + ) + + self.nvfp4_backend = nvfp4_backend_value local_in_features = in_features local_out_features = out_features From a70c06de1d0e3d8e68e0c1f1e97fea4adcb3b978 Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Mon, 1 Dec 2025 08:29:03 +0000 Subject: [PATCH 19/21] clean code Signed-off-by: Shijie Wang --- .../compilation/patterns/ar_residual_norm.py | 30 ++++-------- .../dense_blockscaled_gemm_persistent.py | 48 +++++++++---------- 2 files changed, 33 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py index 34919d86bfdc..bc42e7614941 100644 --- a/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py +++ b/tensorrt_llm/_torch/compilation/patterns/ar_residual_norm.py @@ -604,38 +604,26 @@ def target_nvfp4_gemm_prologue_pattern( def extra_check(match: Match) -> bool: # Validate backend value - backend_node = match.kwargs.get('backend') - if backend_node is None: + backend_value = match.kwargs.get('backend') + if backend_value is None: # No backend specified, use default - OK return True - valid_backends = {'auto', 'cutlass', 'cublaslt', 'cutedsl'} - - # Case 1: backend is a Node with metadata - if hasattr(backend_node, 'meta') and 'val' in backend_node.meta: - backend_value = backend_node.meta['val'] - if isinstance(backend_value, str): - return backend_value in valid_backends - return False # Invalid type - - # Case 2: backend is a constant Node in the graph - if hasattr(backend_node, 'target'): - return backend_node.target in valid_backends - - # Case 3: backend is a Python literal in kwargs - if isinstance(backend_node, str): - return backend_node in valid_backends + # backend should be a string literal + if not isinstance(backend_value, str): + return False - # Unknown format - reject to be safe - return False + valid_backends = {'auto', 'cutlass', 'cublaslt', 'cutedsl'} + return backend_value in valid_backends register_replacement( empty_nvfp4_gemm_prologue_pattern, target_nvfp4_gemm_prologue_pattern, - [extra_check], + [], fwd_only, custom_pass, search_fn_pattern=ub_copy, + extra_check=extra_check, ) def register_mm_prologue(custom_pass: PatternMatcherPass): diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py index b27b096a820e..0b01692d2114 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_persistent.py @@ -548,30 +548,30 @@ class SharedStorage: # GPU device kernel @cute.kernel def kernel( - self, - tiled_mma: cute.TiledMma, - tiled_mma_sfb: cute.TiledMma, - tma_atom_a: cute.CopyAtom, - mA_mkl: cute.Tensor, - tma_atom_b: cute.CopyAtom, - mB_nkl: cute.Tensor, - tma_atom_sfa: cute.CopyAtom, - mSFA_mkl: cute.Tensor, - tma_atom_sfb: cute.CopyAtom, - mSFB_nkl: cute.Tensor, - tma_atom_c: Optional[cute.CopyAtom], - mC_mnl: cute.Tensor, - cluster_layout_vmnk: cute.Layout, - cluster_layout_sfb_vmnk: cute.Layout, - a_smem_layout_staged: cute.ComposedLayout, - b_smem_layout_staged: cute.ComposedLayout, - sfa_smem_layout_staged: cute.Layout, - sfb_smem_layout_staged: cute.Layout, - c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], - epi_tile: cute.Tile, - tile_sched_params: utils.PersistentTileSchedulerParams, - epilogue_op: cutlass.Constexpr, - alpha: cute.Tensor, # Single-element tensor containing alpha value + self, + tiled_mma: cute.TiledMma, + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: Optional[cute.CopyAtom], + mC_mnl: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + alpha: cute.Tensor, ): """ GPU device kernel performing the Persistent batched GEMM computation. From f013872590029d1b1c4d12da7a562eae5f175e9c Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Mon, 1 Dec 2025 11:00:21 +0000 Subject: [PATCH 20/21] minor change Signed-off-by: Shijie Wang --- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 1fb0d7f227cd..d40c1fd58445 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -757,11 +757,11 @@ def get_valid_tactics(self, inputs: List[torch.Tensor], if IS_CUTLASS_DSL_AVAILABLE: # Check SM version first - CuteDSL NVFP4 only supports SM 100 (B200) sm_version = get_sm_version() - if sm_version != 100: + if sm_version not in [100, 103]: if backend == "cutedsl": # Explicitly requested CuteDSL but SM version not supported raise ValueError( - f"CuteDSL NVFP4 backend requires SM 100 (B200), but got SM {sm_version}. " + f"CuteDSL NVFP4 backend requires SM 100 (B200) or SM 103 (B300), but got SM {sm_version}. " f"CuteDSL NVFP4 is not supported on this GPU architecture. " f"Please use backend='auto' to automatically select a compatible backend." ) From 906391b37ae9dabe2a1add288bc98507d7bdbb1f Mon Sep 17 00:00:00 2001 From: Shijie Wang Date: Mon, 1 Dec 2025 11:51:47 +0000 Subject: [PATCH 21/21] fix code style Signed-off-by: Shijie Wang --- tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py | 4 ---- 1 file changed, 4 deletions(-) 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 187fd11fb80f..e0c50eccc3de 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -98,10 +98,6 @@ def get_valid_tactics( a_major = "k" b_major = "k" - # Data types - ab_dtype = cutlass.Float4E2M1FN - c_dtype = cutlass.BFloat16 - # Early exit: Check K dimension alignment # For K-major layout (A and B tensors), K is the major mode (contiguous dimension). # 16-byte alignment requirement: K must be divisible by 32 for FP4 (128 bits / 4 bits = 32)