From 73ca51746fb90a9bc6996fcbab71f3ea5718a040 Mon Sep 17 00:00:00 2001 From: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> Date: Thu, 16 Jul 2026 23:56:36 -0700 Subject: [PATCH 1/2] [None][perf] Fuse DeepSeek-V4 Indexer Q projection with CuTe DSL Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> --- .../fp8_blockscale_quant_packed.cu | 95 +++-- .../fp8_blockscale_quant_packed.h | 6 + cpp/tensorrt_llm/thop/fp8Quantize.cpp | 34 ++ .../sparse/deepseek_v4/deepseek_v4.py | 64 ++- .../_torch/custom_ops/cpp_custom_ops.py | 9 + .../_torch/custom_ops/cute_dsl_custom_ops.py | 314 ++++++++++++++ .../dense_blockscaled_gemm_act_fusion.py | 278 +++++++++++- .../dense_blockscaled_gemm_persistent.py | 402 +++++++++++++++++- tensorrt_llm/_torch/modules/linear.py | 35 +- .../test_lists/test-db/l0_dgx_b200.yml | 1 + .../attention/sparse/test_cpp_custom_ops.py | 102 ++++- 11 files changed, 1280 insertions(+), 60 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu index a4923fdbd072..adaf60ccc3a8 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.cu @@ -52,7 +52,7 @@ __device__ __forceinline__ float reciprocal_approximate_ftz_local(float a) // (8 lanes × 16 BF16 elems = 128 elems). After per-block amax, lanes // 0/8/16/24 each hold one UE8M0 scale byte; lane 0 packs them into a uint32 // and stores in the deep_gemm-expected MN-major layout. -template +template __global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict__ fp8_output, int32_t* __restrict__ packed_scale_output, __nv_bfloat16 const* __restrict__ input, int const m, int const k, int const scale_leading_dim_uint32) @@ -72,6 +72,7 @@ __global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict_ bool const row_in_range = (m_idx < m); uint32_t packed = 0u; + uint32_t scale_byte = 0u; if (row_in_range) { int const k_base = packed_sf_k_idx * 512 + lane_id * 16; @@ -120,6 +121,7 @@ __global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict_ float const dequant_scale_raw = amax * reciprocal_approximate_ftz_local(448.0f); __nv_fp8_e8m0 ue8m0_scale; ue8m0_scale.__x = __nv_cvt_float_to_e8m0(dequant_scale_raw, __NV_SATFINITE, cudaRoundPosInf); + scale_byte = static_cast(ue8m0_scale.__x); // Recover quant_scale = 1 / 2^(exp - 127) for fp8 conversion. constexpr uint32_t FP32_EXPONENT_BIAS = 127u; @@ -162,32 +164,58 @@ __global__ void fp8_quantize_1x128_packed_kernel_impl(__nv_fp8_e4m3* __restrict_ } // ---- 5. Pack 4 UE8M0 scales (lanes 0/8/16/24). ---- - uint32_t const s0 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 0); - uint32_t const s1 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 8); - uint32_t const s2 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 16); - uint32_t const s3 = __shfl_sync(0xFFFFFFFFu, static_cast(ue8m0_scale.__x), 24); - if (lane_id == 0) + if constexpr (!OutputCuteDslSf) { - // Mask off scale bytes whose sf_k is past the actual K. - int const num_sf_k = (k + 127) / 128; - int const sf_k_base = packed_sf_k_idx * 4; - if (sf_k_base + 0 < num_sf_k) - packed |= s0; - if (sf_k_base + 1 < num_sf_k) - packed |= (s1 << 8); - if (sf_k_base + 2 < num_sf_k) - packed |= (s2 << 16); - if (sf_k_base + 3 < num_sf_k) - packed |= (s3 << 24); + uint32_t const s0 = __shfl_sync(0xFFFFFFFFu, scale_byte, 0); + uint32_t const s1 = __shfl_sync(0xFFFFFFFFu, scale_byte, 8); + uint32_t const s2 = __shfl_sync(0xFFFFFFFFu, scale_byte, 16); + uint32_t const s3 = __shfl_sync(0xFFFFFFFFu, scale_byte, 24); + if (lane_id == 0) + { + // Mask off scale bytes whose sf_k is past the actual K. + int const num_sf_k = (k + 127) / 128; + int const sf_k_base = packed_sf_k_idx * 4; + if (sf_k_base + 0 < num_sf_k) + packed |= s0; + if (sf_k_base + 1 < num_sf_k) + packed |= (s1 << 8); + if (sf_k_base + 2 < num_sf_k) + packed |= (s2 << 16); + if (sf_k_base + 3 < num_sf_k) + packed |= (s3 << 24); + } } } - // Always write the packed scale — `packed` is 0 for padded rows. The grid - // covers the full [0, scale_leading_dim_uint32) leading dim (rounded up to - // WarpsPerBlock), and the m_idx guard drops the few rows past the buffer end. - if (lane_id == 0 && m_idx < scale_leading_dim_uint32) + if constexpr (OutputCuteDslSf) + { + // Native MXF8 MMA consumes one UE8M0 scale per 32 K values. Preserve + // the production 1x128 quantization contract by replicating each scale + // four times directly into CUTLASS/CuTe's 128x4 swizzled layout. + if (lane_id % 8 == 0 && m_idx < scale_leading_dim_uint32) + { + int const sf128_idx = packed_sf_k_idx * 4 + lane_id / 8; + int const num_sf128 = (k + 127) / 128; + if (sf128_idx < num_sf128) + { + int const num_sf32 = (k + 31) / 32; + int const num_k_tiles = (num_sf32 + 3) / 4; + int64_t const dst_offset = static_cast(m_idx / 128) * num_k_tiles * 512 + + static_cast(sf128_idx) * 512 + (m_idx % 32) * 16 + ((m_idx % 128) / 32) * 4; + uint32_t const replicated = row_in_range ? scale_byte * 0x01010101u : 0u; + *reinterpret_cast(reinterpret_cast(packed_scale_output) + dst_offset) = replicated; + } + } + } + else { - packed_scale_output[static_cast(packed_sf_k_idx) * scale_leading_dim_uint32 + m_idx] = packed; + // Always write the packed scale — `packed` is 0 for padded rows. The grid + // covers the full [0, scale_leading_dim_uint32) leading dim (rounded up to + // WarpsPerBlock), and the m_idx guard drops the few rows past the buffer end. + if (lane_id == 0 && m_idx < scale_leading_dim_uint32) + { + packed_scale_output[static_cast(packed_sf_k_idx) * scale_leading_dim_uint32 + m_idx] = packed; + } } #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) @@ -215,8 +243,27 @@ void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32 dim3 const block(kWarpsPerBlock * 32, 1, 1); tensorrt_llm::common::launchWithPdlWhenEnabled("fp8_quantize_1x128_packed_kernel_impl", - fp8_quantize_1x128_packed_kernel_impl, grid, block, 0, stream, fp8_output, packed_scale_output, - input, m, k, scale_leading_dim_uint32); + fp8_quantize_1x128_packed_kernel_impl, grid, block, 0, stream, fp8_output, + packed_scale_output, input, m, k, scale_leading_dim_uint32); +} + +void launch_fp8_quantize_1x128_cutedsl_bf16_e4m3(__nv_fp8_e4m3* fp8_output, uint8_t* swizzled_scale_output, + __nv_bfloat16 const* input, int m, int k, int padded_m, cudaStream_t stream) +{ + if (m <= 0 || k <= 0) + { + return; + } + + constexpr int kWarpsPerBlock = 4; + int const num_packed_sf_k = (((k + 127) / 128) + 3) / 4; + int const m_blocks = (padded_m + kWarpsPerBlock - 1) / kWarpsPerBlock; + dim3 const grid(num_packed_sf_k, m_blocks, 1); + dim3 const block(kWarpsPerBlock * 32, 1, 1); + + tensorrt_llm::common::launchWithPdlWhenEnabled("fp8_quantize_1x128_cutedsl_kernel_impl", + fp8_quantize_1x128_packed_kernel_impl, grid, block, 0, stream, fp8_output, + reinterpret_cast(swizzled_scale_output), input, m, k, padded_m); } } // namespace kernels::fp8_blockscale_gemm diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h index a4079cd9b054..7f8c4f07d33d 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/fp8_blockscale_gemm/fp8_blockscale_quant_packed.h @@ -47,6 +47,12 @@ namespace kernels::fp8_blockscale_gemm void launch_fp8_quantize_1x128_packed_bf16_e4m3(__nv_fp8_e4m3* fp8_output, int32_t* packed_scale_output, __nv_bfloat16 const* input, int m, int k, int scale_leading_dim_uint32, cudaStream_t stream); +// Quantizes with the same 1x128 scale as above, but replicates each UE8M0 +// scale over four 32-wide groups and writes the native SM100 CuTe/CUTLASS +// 128x4 swizzled scale layout. +void launch_fp8_quantize_1x128_cutedsl_bf16_e4m3(__nv_fp8_e4m3* fp8_output, uint8_t* swizzled_scale_output, + __nv_bfloat16 const* input, int m, int k, int padded_m, cudaStream_t stream); + } // namespace kernels::fp8_blockscale_gemm TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/fp8Quantize.cpp b/cpp/tensorrt_llm/thop/fp8Quantize.cpp index 35ba5c440e18..43eea8cff838 100644 --- a/cpp/tensorrt_llm/thop/fp8Quantize.cpp +++ b/cpp/tensorrt_llm/thop/fp8Quantize.cpp @@ -209,6 +209,38 @@ std::tuple fp8_quantize_1x128_packed_ue8m0(at::Tensor co return {valueE4M3.slice(0, 0, m), packedScale}; } + +std::tuple fp8_quantize_1x128_cutedsl_ue8m0(at::Tensor const& self) +{ + CHECK_TH_CUDA(self); + CHECK_CONTIGUOUS(self); + + TORCH_CHECK(self.scalar_type() == at::ScalarType::BFloat16, "Input matrix dtype must be BF16."); + TORCH_CHECK(self.dim() == 2, "input must be a matrix"); + TORCH_CHECK(tensorrt_llm::common::isSM100Family(), + "fp8_quantize_1x128_cutedsl_ue8m0 currently only supports SM100 (Blackwell)."); + + auto const m = self.sizes()[0]; + auto const k = self.sizes()[1]; + TORCH_CHECK(m <= std::numeric_limits::max(), "M must be within int32"); + TORCH_CHECK(k <= std::numeric_limits::max(), "K must be within int32"); + TORCH_CHECK(k % 128 == 0, "K must be divisible by the production FP8 block size 128, but got ", k); + + at::Tensor valueE4M3 + = at::detail::empty_cuda({m, k}, at::ScalarType::Float8_e4m3fn, self.device(), /* stride */ std::nullopt); + auto const paddedM = (m + 127) / 128 * 128; + auto const sfCols = (k / 32 + 3) / 4 * 4; + at::Tensor scaleE8M0 + = at::detail::empty_cuda({paddedM * sfCols}, at::ScalarType::Byte, self.device(), /* stride */ std::nullopt); + + auto stream = at::cuda::getCurrentCUDAStream(self.get_device()); + tensorrt_llm::kernels::fp8_blockscale_gemm::launch_fp8_quantize_1x128_cutedsl_bf16_e4m3( + reinterpret_cast<__nv_fp8_e4m3*>(valueE4M3.data_ptr()), scaleE8M0.data_ptr(), + reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), static_cast(m), static_cast(k), + static_cast(paddedM), stream); + + return {valueE4M3, scaleE8M0}; +} } // namespace torch_ext TRTLLM_NAMESPACE_END @@ -218,6 +250,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def("fp8_quantize_1x128(Tensor input, bool use_ue8m0=False) -> (Tensor, Tensor)"); m.def("fp8_batched_quantize_1x128_permute102(Tensor input) -> (Tensor, Tensor)"); m.def("fp8_quantize_1x128_packed_ue8m0(Tensor input) -> (Tensor, Tensor)"); + m.def("fp8_quantize_1x128_cutedsl_ue8m0(Tensor input) -> (Tensor, Tensor)"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) @@ -225,4 +258,5 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m) m.impl("fp8_quantize_1x128", &tensorrt_llm::torch_ext::fp8_quantize_1x128); m.impl("fp8_batched_quantize_1x128_permute102", &tensorrt_llm::torch_ext::fp8_batched_quantize_1x128_permute102); m.impl("fp8_quantize_1x128_packed_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_packed_ue8m0); + m.impl("fp8_quantize_1x128_cutedsl_ue8m0", &tensorrt_llm::torch_ext::fp8_quantize_1x128_cutedsl_ue8m0); } diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 52f16c84ffe9..c7ef7c19eb23 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -32,7 +32,7 @@ from tensorrt_llm._torch.modules.multi_stream_utils import do_multi_stream from tensorrt_llm._torch.modules.rotary_embedding import RotaryEmbedding from tensorrt_llm._torch.utils import maybe_compile -from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm._utils import is_sm_100f, prefer_pinned from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.utils import fp8_utils from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole @@ -1040,6 +1040,11 @@ def __init__( layer_idx, aux_stream, ) + # Preserve the checkpoint's 128x128 FP8 quantization while deriving a + # native-MXF8 (sf_vec=32) scale view for the TRT-LLM CuTe DSL fused + # indexer-Q projection. FP8BlockScalesLinearMethod materializes the + # derived, swizzled scale once after weight loading. + self.wq_b.use_indexer_q_cutedsl_fusion = True # Override base Indexer.weights_proj to bf16 (matches V4 checkpoint). self.weights_proj = Linear( self.hidden_size, @@ -1099,6 +1104,10 @@ def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): """ q = self.wq_b(qr) q = q.view(-1, self.n_heads, self.head_dim) + return self._apply_q_rope(q, position_ids) + + def _apply_q_rope(self, q: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + """Apply RoPE in-place to a projected indexer Q tensor.""" # Fused in-place RoPE on the rope portion of each head nope_dim = self.head_dim - self.rope_dim torch.ops.trtllm.mla_rope_inplace( @@ -1113,6 +1122,49 @@ def _qk_projection_and_rope(self, qr: torch.Tensor, position_ids: torch.Tensor): ) return q + def _project_and_quantize_q( + self, qr: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Project Q and produce the cache precision consumed by the indexer. + + The DSv4 MXFP4 configuration can fuse projection, interleaved RoPE, + and FP4 quantization. If the optional Hadamard transform is active, + retain the legacy path because that transform changes all 128 values + in a head and is not part of the CuTe DSL kernel. + """ + use_fused_project_mxfp4 = ( + self.indexer_cache_dtype == KVCacheDtype.MXFP4_BLOCKWISE + and not HAS_FAST_HADAMARD + and not self.rotary_emb.is_neox + and self.head_dim == 128 + and self.rope_dim == 64 + and self.wq_b.has_fp8_block_scales + and hasattr(self.wq_b, "indexer_q_weight_scale_cutedsl") + and hasattr( + torch.ops.trtllm, + "cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", + ) + and qr.dtype == torch.bfloat16 + and is_sm_100f() + ) + if use_fused_project_mxfp4: + q_fp4, q_scale = torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( + qr, + self.wq_b.weight, + self.wq_b.indexer_q_weight_scale_cutedsl, + position_ids.view(-1), + self.rotary_emb.rotary_cos_sin.view(-1, self.rope_dim), + self.wq_b.indexer_q_alpha_cutedsl, + use_tvm_ffi=True, + ) + return q_fp4.view(-1, self.n_heads, self.head_dim // 2), q_scale.view( + -1, self.n_heads, 1 + ) + + q = self.wq_b(qr).view(-1, self.n_heads, self.head_dim) + q = self._apply_q_rope(q, position_ids) + return self._quantize_q(q) + def _quantize_q(self, q: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: # Rotate + quantize (layout matches compressor K: [nope|pe]). After # rotate_activation (Hadamard) the nope/rope split becomes a linear @@ -1242,7 +1294,7 @@ def _run_overlapped_indexer_prepare( if pre_aux is None: self.indexer_start_event.record() - q = self._qk_projection_and_rope(qr, position_ids) + q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) with torch.cuda.stream(self.aux_stream): self.indexer_start_event.wait() @@ -1263,9 +1315,7 @@ def _run_overlapped_indexer_prepare( k_fp8.record_stream(cur_stream) if k_scale is not None: k_scale.record_stream(cur_stream) - q = self._qk_projection_and_rope(qr, position_ids) - - q_fp8, q_scale = self._quantize_q(q) + q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) self.weights_proj_event.wait() weights = self._apply_weight_scale(weights, q_scale) @@ -1282,9 +1332,7 @@ def _run_serial_indexer_prepare( ) -> Tuple[ torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], torch.Tensor ]: - q = self._qk_projection_and_rope(qr, position_ids) - - q_fp8, q_scale = self._quantize_q(q) + q_fp8, q_scale = self._project_and_quantize_q(qr, position_ids) weights = self.weights_proj(hidden_states) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index e86993ef8f67..cf7240b3be36 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -731,6 +731,15 @@ def _(input: torch.Tensor): dtype=torch.float8_e4m3fn), input.new_empty( (m, num_packed_sf_k), dtype=torch.int32) + @torch.library.register_fake("trtllm::fp8_quantize_1x128_cutedsl_ue8m0") + def _(input: torch.Tensor): + m, k = input.shape + padded_m = fp4_utils.pad_up(m, 128) + sf_cols = fp4_utils.pad_up(k // 32, 4) + return torch.empty_like(input, + dtype=torch.float8_e4m3fn), input.new_empty( + (padded_m * sf_cols, ), dtype=torch.uint8) + @torch.library.register_fake("trtllm::causal_conv1d_fwd") def _( x: torch.Tensor, 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 0f674cb20e67..e6bae5c80c42 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -3636,6 +3636,320 @@ def _fake_single_b( device=input_scale.device) return output, output_scale + _INDEXER_Q_CUTEDSL_TUNING_BUCKETS = ( + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + ) + + def _map_cutedsl_indexer_q_tuning_bucket(num_tokens: int) -> int: + if num_tokens <= 4: + return 4 + if num_tokens <= 8: + return 8 + if num_tokens <= 16: + return 16 + return max(32, last_positive_power_of_2(num_tokens)) + + def _prepare_cutedsl_indexer_q_tuning_inputs( + inputs: List[torch.Tensor]) -> List[torch.Tensor]: + inputs[3] = torch.zeros_like(inputs[3]) + return inputs + + class CuteDSLIndexerQBlackwellRunner(TunableRunner): + """Native MXF8 GEMM with fused DSv4 indexer-Q RoPE/MXFP4 output.""" + + kernel_class = Sm100BlockScaledPersistentDenseGemmActFusionKernel + small_m_kernel_class = Sm100BlockScaledPersistentDenseGemmKernel + kernel_cache = dict() + tuning_config = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, + 0, + _INDEXER_Q_CUTEDSL_TUNING_BUCKETS, + _map_cutedsl_indexer_q_tuning_bucket, + ), ), + constraint_specs=(ConstraintSpec(3, 0, + lambda shapes: shapes[0][0]), ), + inputs_pre_hook=_prepare_cutedsl_indexer_q_tuning_inputs, + use_cold_l2_cache=True, + # CuTe kernels are JIT compiled into a process-local cache while + # the autotuner profiles tactics. Never persist only the selected + # tactic: a new process would otherwise skip that compilation and + # pay for cute.compile in its first inference forward. + exclude_from_cache=True, + # Every rank owns a process-local CuTe module cache. Profiling in + # parallel across ranks would leave the winning tactic uncompiled + # on ranks that benchmarked a different subset. + distributed_tuning_strategy=DistributedTuningStrategy.INDEPENDENT, + ) + + _small_m_tactics = ( + ("swap_ab", (128, 16), (1, 1), False, 4), + ("swap_ab", (128, 16), (1, 1), False, 8), + ) + _native_tactics = ( + ("native", (128, 128), (1, 1), False, 0), + ("native", (128, 128), (1, 2), False, 0), + ("native", (128, 128), (2, 1), False, 0), + ("native", (128, 128), (2, 1), True, 0), + ("native", (256, 128), (2, 1), False, 0), + ("native", (256, 128), (2, 1), True, 0), + ("native", (256, 128), (2, 2), True, 0), + ) + + def __init__(self, use_tvm_ffi: bool = True): + super().__init__() + self.use_tvm_ffi = use_tvm_ffi + + def unique_id(self): + return (self.use_tvm_ffi, ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple]: + if not is_sm_100f(): + return [] + m, k = inputs[0].shape + n = inputs[1].shape[0] + tactics = [] + if 0 < m <= 16 and n % 128 == 0 and k % 128 == 0: + tactics.extend(self.__class__._small_m_tactics) + + tactics.extend([ + tactic for tactic in self.__class__._native_tactics + if self.__class__.kernel_class.can_implement( + cutlass.Float8E4M3FN, + cutlass.Float8E8M0FNU, + 32, + cutlass.Float4E2M1FN, + tactic[1], + tactic[2], + m, + n, + k, + 1, + "k", + "k", + "n", + ) + ]) + return tactics + + @staticmethod + def _fallback_tactic(m: int) -> Tuple: + """Safe eager-mode fallback when TRT-LLM autotuning is disabled.""" + if m <= 4: + return ("swap_ab", (128, 16), (1, 1), False, 4) + if m <= 8: + return ("swap_ab", (128, 16), (1, 1), False, 8) + return ("native", (256, 128), (2, 1), False, 0) + + @staticmethod + def _ptr(tensor: torch.Tensor, dtype, align: int = 16): + return make_ptr( + dtype, + tensor.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=align, + ) + + def forward( + self, + inputs: List[torch.Tensor], + tactic, + ) -> Tuple[torch.Tensor, torch.Tensor]: + input, weight, weight_scale, position_ids, cos_sin_cache, alpha = inputs + m, k = input.shape + n = weight.shape[0] + if tactic == -1: + tactic = self._fallback_tactic(m) + (kernel_kind, mma_tiler_mn, cluster_shape_mn, use_prefetch, + transform_warps) = tactic + if kernel_kind == "swap_ab": + if not 0 < m <= mma_tiler_mn[1]: + raise ValueError( + "The small-M indexer-Q kernel requires one non-empty " + f"token tile: M={m}, tile N={mma_tiler_mn[1]}") + if n % 128 != 0 or k % 128 != 0: + raise ValueError( + "The small-M indexer-Q kernel requires N and K to be " + f"divisible by 128, but got N={n}, K={k}") + + a, a_sf = torch.ops.trtllm.fp8_quantize_1x128_cutedsl_ue8m0(input) + packed = torch.empty((m, n // 2), + dtype=torch.uint8, + device=input.device) + output_scale = torch.empty((m, n // 32), + dtype=torch.uint8, + device=input.device) + + a_ptr = self._ptr(a, cutlass.Float8E4M3FN) + b_ptr = self._ptr(weight, cutlass.Float8E4M3FN) + a_sf_ptr = self._ptr(a_sf, cutlass.Float8E8M0FNU) + b_sf_ptr = self._ptr(weight_scale, cutlass.Float8E8M0FNU) + packed_ptr = self._ptr(packed, cutlass.Uint8) + output_scale_ptr = self._ptr(output_scale, cutlass.Float8E8M0FNU) + position_ids_ptr = self._ptr(position_ids, cutlass.Int32, 4) + cos_sin_ptr = self._ptr(cos_sin_cache, cutlass.Float32, 32) + alpha_cute = cute.runtime.from_dlpack(alpha) + + if self.use_tvm_ffi: + stream = cute.runtime.make_fake_stream( + use_tvm_ffi_env_stream=True) + else: + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + cache_key = (kernel_kind, mma_tiler_mn, cluster_shape_mn, + use_prefetch, transform_warps, self.use_tvm_ffi) + if cache_key not in self.__class__.kernel_cache: + if kernel_kind == "swap_ab": + gemm = self.__class__.small_m_kernel_class( + 32, + mma_tiler_mn, + cluster_shape_mn, + use_prefetch=use_prefetch, + indexer_q_fusion=True, + indexer_transform_warps=transform_warps, + ) + compile_entry = gemm.wrapper_indexer_q_swap_ab + else: + gemm = self.__class__.kernel_class( + 32, + mma_tiler_mn, + cluster_shape_mn, + True, + use_prefetch, + activation_type=ActivationType.Identity, + indexer_q_fusion=True, + ) + compile_entry = gemm.wrapper_indexer_q + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + compiled = cute.compile( + compile_entry, + m, + n, + k, + pad_up(m, 128) // 128, + pad_up(n, 128) // 128, + pad_up(k // 32, 4) // 4, + cos_sin_cache.shape[0], + 1, + a_ptr, + b_ptr, + a_sf_ptr, + b_sf_ptr, + packed_ptr, + output_scale_ptr, + position_ids_ptr, + cos_sin_ptr, + alpha_cute, + max_active_clusters, + stream, + options="--opt-level 2 --enable-tvm-ffi" + if self.use_tvm_ffi else "--opt-level 2", + ) + self.__class__.kernel_cache[cache_key] = compiled + else: + compiled = self.__class__.kernel_cache[cache_key] + + dynamic_args = [ + m, + n, + k, + pad_up(m, 128) // 128, + pad_up(n, 128) // 128, + pad_up(k // 32, 4) // 4, + cos_sin_cache.shape[0], + ] + if self.use_tvm_ffi: + compiled( + *dynamic_args, + a.data_ptr(), + weight.data_ptr(), + a_sf.data_ptr(), + weight_scale.data_ptr(), + packed.data_ptr(), + output_scale.data_ptr(), + position_ids.data_ptr(), + cos_sin_cache.data_ptr(), + alpha, + ) + else: + compiled( + *dynamic_args, + a_ptr, + b_ptr, + a_sf_ptr, + b_sf_ptr, + packed_ptr, + output_scale_ptr, + position_ids_ptr, + cos_sin_ptr, + alpha_cute, + stream, + ) + return packed.view(torch.int8), output_scale.view(torch.int32) + + @torch.library.custom_op( + "trtllm::cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", + mutates_args=(), + device_types="cuda", + ) + def cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + position_ids: torch.Tensor, + cos_sin_cache: torch.Tensor, + alpha: torch.Tensor, + use_tvm_ffi: bool = True, + ) -> Tuple[torch.Tensor, torch.Tensor]: + runner = CuteDSLIndexerQBlackwellRunner(use_tvm_ffi) + inputs = [ + input, weight, weight_scale, position_ids, cos_sin_cache, alpha + ] + tuner = AutoTuner.get() + _, tactic = tuner.choose_one( + "trtllm::cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell", + [runner], + runner.__class__.tuning_config, + inputs, + ) + return runner(inputs, tactic=tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell") + def _( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + position_ids: torch.Tensor, + cos_sin_cache: torch.Tensor, + alpha: torch.Tensor, + use_tvm_ffi: bool = True, + ): + m, n = input.shape[0], weight.shape[0] + return ( + input.new_empty((m, n // 2), dtype=torch.int8), + input.new_empty((m, n // 128), dtype=torch.int32), + ) + class CuteDSLFp8BlackwellRunner(TunableRunner): kernel_class = Sm100BlockwiseGemmKernel kernel_cache = dict() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py index b96186049f67..88215a759fc2 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/dense_blockscaled_gemm_act_fusion.py @@ -125,6 +125,7 @@ def __init__( vectorized_f32: bool, use_prefetch: bool = False, activation_type: ActivationType = ActivationType.Swiglu, + indexer_q_fusion: bool = False, ): """Initializes the configuration for a Blackwell dense GEMM kernel with fused activation. @@ -182,19 +183,28 @@ def __init__( self.vectorized_f32 = vectorized_f32 self.activation_type = activation_type + self.indexer_q_fusion = indexer_q_fusion self.is_gated = is_gated_activation(activation_type) # Precompute per-activation flags as plain Python bools so the epilogue # dispatch can use cutlass.const_expr(...) on them (like is_gated). An # inline enum comparison inside @cute.jit is not folded as a constant. self._act_is_swiglu = activation_type == ActivationType.Swiglu self._act_is_gelu = activation_type == ActivationType.Gelu + self._act_is_identity = activation_type == ActivationType.Identity # Host-side guard (raise is not allowed inside the @cute.kernel epilogue): # only SwiGLU and GELU(tanh) are wired in the fused epilogue. - if not (self._act_is_swiglu or self._act_is_gelu): + if not (self._act_is_swiglu or self._act_is_gelu or self._act_is_identity): raise NotImplementedError( f"Fused epilogue activation {activation_type} not implemented " - f"(only Swiglu and Gelu are wired)." + f"(only Swiglu, Gelu, and the indexer-Q identity epilogue are wired)." ) + if self.indexer_q_fusion: + if activation_type != ActivationType.Identity: + raise ValueError("indexer_q_fusion requires ActivationType.Identity") + if sf_vec_size != 32 or mma_tiler_mn[1] != 128: + raise ValueError( + "indexer_q_fusion requires MXF8 sf_vec_size=32 and a 128-column MMA tile" + ) def _setup_attributes(self): """Set up configurations that are dependent on GEMM inputs @@ -298,7 +308,7 @@ def _setup_attributes(self): self.is_b_mcast = self.num_mcast_ctas_b > 1 self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 - # SwiGLU: hardcoded epilogue tile matching grouped swiglu variant + # SwiGLU: hardcoded epilogue tile matching grouped swiglu variant. self.epi_tile = (128, 64) self.epi_tile_cnt = ( self.cta_tile_shape_mnk_c[0] // self.epi_tile[0], @@ -389,6 +399,9 @@ def __call__( sfc_tensor: Optional[cute.Tensor] = None, norm_const_tensor: Optional[cute.Tensor] = None, bias_tensor: Optional[cute.Tensor] = None, + indexer_scale_tensor: Optional[cute.Tensor] = None, + position_ids_tensor: Optional[cute.Tensor] = None, + cos_sin_cache_tensor: Optional[cute.Tensor] = None, ): """Execute the GEMM operation with SwiGLU fusion in steps: - Setup static attributes before smem/grid/tma computation @@ -444,6 +457,20 @@ def __call__( # Setup sfc tensor by filling C tensor to scale factor atom layout self.generate_sfc = sfc_tensor is not None and norm_const_tensor is not None + if cutlass.const_expr( + self.indexer_q_fusion + and ( + not self.generate_sfc + or self.c_dtype != cutlass.Float4E2M1FN + or indexer_scale_tensor is None + or position_ids_tensor is None + or cos_sin_cache_tensor is None + ) + ): + raise ValueError( + "indexer_q_fusion requires FP4 C, SFC generation, contiguous " + "output scales, position IDs, and a cos/sin cache" + ) if cutlass.const_expr(self.generate_sfc): sfc_layout = blockscaled_utils.tile_atom_to_shape_SF(c_tensor.shape, self.sf_vec_size) sfc_tensor = cute.make_tensor(sfc_tensor.iterator, sfc_layout) @@ -621,6 +648,9 @@ class SharedStorage: sfc_tensor, norm_const_tensor, bias_tensor, + indexer_scale_tensor, + position_ids_tensor, + cos_sin_cache_tensor, self.cluster_layout_vmnk, self.cluster_layout_sfb_vmnk, self.a_smem_layout_staged, @@ -662,6 +692,9 @@ def kernel( mSFC_mnl: Optional[cute.Tensor], norm_const_tensor: Optional[cute.Tensor], mBias_mnl: Optional[cute.Tensor], + mIndexerScale: Optional[cute.Tensor], + mPositionIds: Optional[cute.Tensor], + mCosSinCache: Optional[cute.Tensor], cluster_layout_vmnk: cute.Layout, cluster_layout_sfb_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, @@ -1365,6 +1398,7 @@ def kernel( tTR_rAcc_up, tTR_rAcc_gate, tTR_gBias_base, + tTR_cC, ) = self.epilog_tmem_copy_and_partition( epi_tidx, tCtAcc_base, tCgC, epi_tile, use_2cta_instrs, tCgBias ) @@ -1503,6 +1537,7 @@ def kernel( # up * silu(gate) # subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + packed_indexer_scale = cutlass.Uint32(0) for subtile_idx in cutlass.range(0, subtile_cnt, 2 if self.is_gated else 1): if cutlass.const_expr(self.is_gated): @@ -1568,6 +1603,60 @@ def kernel( self._apply_gelu_epilogue( acc_vec_up, alpha_val, tCompute, bias_vec=bias_vec ) + elif cutlass.const_expr(self._act_is_identity): + # Preserve the existing Q-projection contract: the GEMM + # accumulator is rounded to BF16 before RoPE. The + # indexer-Q block below performs that boundary rounding; + # this copy keeps the accumulator in FP32 until then. + tCompute.store(acc_vec_up) + + if cutlass.const_expr(self.indexer_q_fusion): + # The legacy path rounds GEMM output to BF16 before + # RoPE, then rounds the rotated values to BF16 again. + tCompute.store(tCompute.load().to(cutlass.BFloat16).to(cutlass.Float32)) + subtile_n_for_rope = real_subtile_idx % self.epi_tile_cnt[1] + subtiles_per_head = 128 // self.epi_tile[1] + rotary_subtile_begin = 64 // self.epi_tile[1] + subtile_in_head = subtile_n_for_rope % subtiles_per_head + if subtile_in_head >= rotary_subtile_begin: + subtile_m_for_rope = real_subtile_idx // self.epi_tile_cnt[1] + rope_base_m = ( + cur_tile_coord[0] * self.cta_tile_shape_mnk_c[0] + + subtile_m_for_rope * self.epi_tile[0] + ) + rope_m_idx = rope_base_m + tTR_cC[0][0] + if rope_m_idx < mPositionIds.shape[0]: + position = mPositionIds[rope_m_idx] + rope_row = mCosSinCache[position, None] + copy_atom_f32x8 = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=256, + ) + cos_values = cute.make_rmem_tensor((8,), cutlass.Float32) + sin_values = cute.make_rmem_tensor((8,), cutlass.Float32) + for chunk_idx in cutlass.range_constexpr(4): + cos_tile = cute.local_tile(rope_row, (8,), (chunk_idx,)) + sin_tile = cute.local_tile(rope_row, (8,), (chunk_idx + 4,)) + cute.copy( + copy_atom_f32x8, + cute.coalesce(cos_tile), + cute.coalesce(cos_values), + ) + cute.copy( + copy_atom_f32x8, + cute.coalesce(sin_tile), + cute.coalesce(sin_values), + ) + for pair_in_chunk in cutlass.range_constexpr(8): + pair_idx = chunk_idx * 8 + pair_in_chunk + cosine = cos_values[pair_in_chunk] + sine = sin_values[pair_in_chunk] + x = tCompute[pair_idx * 2] + y = tCompute[pair_idx * 2 + 1] + tCompute[pair_idx * 2] = cosine * x - sine * y + tCompute[pair_idx * 2 + 1] = cosine * y + sine * x + tCompute.store(tCompute.load().to(cutlass.BFloat16).to(cutlass.Float32)) if cutlass.const_expr(self.generate_sfc): # @@ -1639,19 +1728,52 @@ def kernel( * norm_const ) - # TODO: need to add f32x2 -> f8x2 conversion - tCrSFC.store(tCrSFC_pvscale.load().to(self.sf_dtype)) + # CuTe's packed f32x2 -> UE8M0 conversion currently + # lowers a two-element result through vector<0xi32>. + # MXF8 with sf_vec_size=32 has exactly two output scale + # values per epilogue thread, so keep these conversions + # scalar until that lowering is fixed upstream. + if cutlass.const_expr(self.sf_vec_size == 32): + for vi in cutlass.range_constexpr(cute.size(tCrSFC)): + tCrSFC[vi] = tCrSFC_pvscale[vi].to(self.sf_dtype) + else: + tCrSFC.store(tCrSFC_pvscale.load().to(self.sf_dtype)) # # Store SFC to global memory # - cute.autovec_copy(tCrSFC, tCgSFC) + if cutlass.const_expr(self.indexer_q_fusion): + subtile_m = real_subtile_idx // self.epi_tile_cnt[1] + subtile_n = real_subtile_idx % self.epi_tile_cnt[1] + scale_m = ( + cur_tile_coord[0] * self.cta_tile_shape_mnk_c[0] + + subtile_m * self.epi_tile[0] + + tTR_cC[0][0] + ) + if scale_m < mIndexerScale.shape[0]: + scale_bytes = cute.recast_tensor(tCrSFC, cutlass.Uint8) + for vi in cutlass.range_constexpr(cute.size(tCrSFC)): + scale_byte_idx = subtile_n * 2 + vi + packed_indexer_scale = packed_indexer_scale | ( + cutlass.Uint32(scale_bytes[vi]) + << cutlass.Uint32(scale_byte_idx * 8) + ) + else: + cute.autovec_copy(tCrSFC, tCgSFC) # # Compute quantized output values and convert to C type # - # TODO: need to add f8x2 -> f32x2 conversion - tCrSFC_qpvscale_up = tCrSFC.load().to(cutlass.Float32) + # Same two-element UE8M0 lowering issue applies in the + # reverse direction. + if cutlass.const_expr(self.sf_vec_size == 32): + tCrSFC_qpvscale_up = cute.make_rmem_tensor( + tCrSFC.shape, cutlass.Float32 + ) + for vi in cutlass.range_constexpr(cute.size(tCrSFC)): + tCrSFC_qpvscale_up[vi] = cutlass.Float32(tCrSFC[vi]) + else: + tCrSFC_qpvscale_up = tCrSFC.load().to(cutlass.Float32) fp32_max = cutlass.Float32(3.40282346638528859812e38) if cutlass.const_expr(self.vectorized_f32): for vi in cutlass.range_constexpr(0, cute.size(tCrSFC), 2): @@ -1686,6 +1808,30 @@ def kernel( acc_vec = tiled_copy_r2s.retile(tCompute).load() tRS_rC.store(acc_vec.to(self.c_dtype)) + if cutlass.const_expr(self.indexer_q_fusion): + packed_words = cute.recast_tensor(tTR_rC, cutlass.Uint32) + for word_idx in cutlass.range_constexpr(cute.size(packed_words)): + packed = packed_words[word_idx] + midpoint_adjust = cutlass.Uint32(0) + for nibble_idx in cutlass.range_constexpr(8): + scaled_value = tCompute[word_idx * 8 + nibble_idx] + abs_value = cute.arch.fmax(scaled_value, -scaled_value) + nibble_adjust = cutlass.Uint32(1) << (nibble_idx * 4) + if abs_value == cutlass.Float32(0.75): + midpoint_adjust = midpoint_adjust | nibble_adjust + if abs_value == cutlass.Float32(1.75): + midpoint_adjust = midpoint_adjust | nibble_adjust + if abs_value == cutlass.Float32(3.5): + midpoint_adjust = midpoint_adjust | nibble_adjust + packed = packed - midpoint_adjust + magnitude = packed & cutlass.Uint32(0x77777777) + nonzero_sign = ( + (magnitude | (magnitude << 1) | (magnitude << 2)) + & cutlass.Uint32(0x44444444) + ) << 1 + packed_words[word_idx] = packed & ( + cutlass.Uint32(0x77777777) | nonzero_sign + ) else: # # Convert to C type (non-SFC path) @@ -1732,6 +1878,10 @@ def kernel( barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads, ) + if cutlass.const_expr(self.indexer_q_fusion and self.generate_sfc): + scale_m = cur_tile_coord[0] * self.cta_tile_shape_mnk_c[0] + tTR_cC[0][0] + if scale_m < mIndexerScale.shape[0]: + mIndexerScale[scale_m, cur_tile_coord[1], 0] = packed_indexer_scale # # Async arrive accumulator buffer empty @@ -1944,7 +2094,14 @@ def epilog_tmem_copy_and_partition( epi_tile: cute.Tile, use_2cta_instrs: Union[cutlass.Boolean, bool], tCgBias: Optional[cute.Tensor] = None, - ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor, Optional[cute.Tensor]]: + ) -> Tuple[ + cute.TiledCopy, + cute.Tensor, + cute.Tensor, + cute.Tensor, + Optional[cute.Tensor], + cute.Tensor, + ]: """ Make tiledCopy for tensor memory load, then use it to partition tensor memory (source) and register array (destination). @@ -2007,6 +2164,8 @@ def epilog_tmem_copy_and_partition( tTR_rAcc_gate = cute.make_rmem_tensor( tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype ) + cC = cute.make_identity_tensor(epi_tile) + tTR_cC = thr_copy_t2r.partition_D(cC) # Partition the per-N bias EXACTLY like C (broadcast over M via stride 0). # Same flat_divide + partition_D path -> tTR_gBias has the same layout as @@ -2018,7 +2177,14 @@ def epilog_tmem_copy_and_partition( ) tTR_gBias = thr_copy_t2r.partition_D(gBias_mnl_epi) - return tiled_copy_t2r, tTR_tAcc, tTR_rAcc_up, tTR_rAcc_gate, tTR_gBias + return ( + tiled_copy_t2r, + tTR_tAcc, + tTR_rAcc_up, + tTR_rAcc_gate, + tTR_gBias, + tTR_cC, + ) def epilog_smem_copy_and_partition( self, @@ -2606,6 +2772,98 @@ def wrapper( bias_tensor=bias_tensor, ) + @cute.jit + def wrapper_indexer_q( + self, + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + sf_m: cutlass.Int64, + sf_n: cutlass.Int64, + sf_k: cutlass.Int64, + cos_sin_rows: cutlass.Int64, + l: cutlass.Constexpr, # noqa: E741 + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + a_sf_ptr: cute.Pointer, + b_sf_ptr: cute.Pointer, + packed_ptr: cute.Pointer, + scale_ptr: cute.Pointer, + position_ids_ptr: cute.Pointer, + cos_sin_ptr: cute.Pointer, + alpha_tensor: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + current_stream: cuda.CUstream, + ): + """MXF8 GEMM with the DeepSeek-V4 indexer-Q RoPE/MXFP4 epilogue. + + ``packed_ptr`` and ``scale_ptr`` use the indexer's existing contiguous + layouts: uint8 [M, N/2] and UE8M0 [M, N/32]. The output tensor is a + logical FP4 view so the repository's native FP4 shared-memory layout + and TMA store path can be reused. + """ + a_tensor = cute.make_tensor( + a_ptr, + layout=cute.make_ordered_layout((m, k, l), order=(1, 0, 2)), + ) + b_tensor = cute.make_tensor( + b_ptr, + layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2)), + ) + c_tensor = cute.make_tensor( + cute.recast_ptr(packed_ptr, dtype=cutlass.Float4E2M1FN), + layout=cute.make_ordered_layout((m, n, l), order=(1, 0, 2)), + ) + scale_tensor = cute.make_tensor( + cute.recast_ptr(scale_ptr, dtype=cutlass.Uint32), + layout=cute.make_ordered_layout((m, n // 128, l), order=(1, 0, 2)), + ) + # generate_sfc needs a scale-factor-shaped tensor to derive its register + # partition. Indexer-Q stores the values through scale_tensor above, + # preserving the consumer's contiguous per-head four-byte layout. + sfc_tensor = cute.make_tensor( + scale_ptr, + layout=cute.make_ordered_layout((32, 4, sf_m, 4, sf_n, l), order=(2, 1, 4, 0, 3, 5)), + ) + position_ids_tensor = cute.make_tensor( + position_ids_ptr, + layout=cute.make_layout((m,)), + ) + cos_sin_cache_tensor = cute.make_tensor( + cos_sin_ptr, + layout=cute.make_ordered_layout((cos_sin_rows, 64), order=(1, 0)), + ) + sfa_tensor = cute.make_tensor( + a_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, sf_m, 4, sf_k, l), + order=(2, 1, 4, 0, 3, 5), + ), + ) + sfb_tensor = cute.make_tensor( + b_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, sf_n, 4, sf_k, l), + order=(2, 1, 4, 0, 3, 5), + ), + ) + + self( + a_tensor, + b_tensor, + sfa_tensor, + sfb_tensor, + c_tensor, + alpha_tensor, + max_active_clusters, + current_stream, + sfc_tensor=sfc_tensor, + norm_const_tensor=alpha_tensor, + indexer_scale_tensor=scale_tensor, + position_ids_tensor=position_ids_tensor, + cos_sin_cache_tensor=cos_sin_cache_tensor, + ) + @cute.jit def wrapper_fp4out( self, 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 78c8532414b7..0194cbc8359e 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 @@ -52,13 +52,48 @@ import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.utils.blockscaled_layout as blockscaled_utils +from cutlass._mlir.dialects import llvm from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import dsl_user_op from .custom_pipeline import PipelineTmaUmma, PipelineUmmaAsync from .utils import (TRTLLM_ENABLE_PDL, griddepcontrol_launch_dependents, griddepcontrol_wait, is_power_of_2) +@dsl_user_op +def _indexer_q_pack_fp4x4(value0: cutlass.Float32, + value1: cutlass.Float32, + value2: cutlass.Float32, + value3: cutlass.Float32, + *, + loc=None, + ip=None) -> cutlass.Uint16: + """Pack four FP32 values with two native E2M1x2 conversions.""" + return cutlass.Uint16( + llvm.inline_asm( + cutlass.Uint16.mlir_type, + [ + value0.ir_value(loc=loc, ip=ip), + value1.ir_value(loc=loc, ip=ip), + value2.ir_value(loc=loc, ip=ip), + value3.ir_value(loc=loc, ip=ip), + ], + """{ + .reg .b8 byte0, byte1; + cvt.rn.satfinite.e2m1x2.f32 byte0, $2, $1; + cvt.rn.satfinite.e2m1x2.f32 byte1, $4, $3; + mov.b16 $0, {byte0, byte1}; + }""", + "=h,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + )) + + class Sm100BlockScaledPersistentDenseGemmKernel: """Implements batched matrix multiplication (C = A x SFA x B x SFB) with support for various data types and Blackwell GPU architectural features, including persistent tile scheduling and warp specialization. @@ -100,6 +135,8 @@ def __init__( use_prefetch: bool = False, swizzle_size: int = 1, raster_along_m: bool = True, + indexer_q_fusion: bool = False, + indexer_transform_warps: int = 4, ): """Initializes the configuration for a Blackwell dense GEMM kernel. @@ -134,6 +171,16 @@ def __init__( self.use_prefetch = use_prefetch self.swizzle_size = swizzle_size self.raster_along_m = raster_along_m + self.indexer_q_fusion = indexer_q_fusion + self.indexer_transform_warps = indexer_transform_warps + if self.indexer_q_fusion and (sf_vec_size != 32 or mma_tiler_mn[0] + != 128 or mma_tiler_mn[1] != 16 + or cluster_shape_mn != (1, 1) + or indexer_transform_warps not in (4, 8)): + raise ValueError( + "The swapped indexer-Q epilogue currently requires " + "MXF8, a 128x16 MMA tile, a 1x1 cluster, and 4 or 8 " + "row-transform warps") self.cta_group = (tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE) @@ -147,8 +194,12 @@ def __init__( ) self.mma_warp_id = 4 self.tma_warp_id = 5 + self.indexer_extra_warp_id = tuple( + range(self.tma_warp_id + 1, self.tma_warp_id + 1 + + indexer_transform_warps - 4)) if self.indexer_q_fusion else () self.threads_per_cta = 32 * len( - (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id)) + (self.mma_warp_id, self.tma_warp_id, *self.epilog_warp_id, + *self.indexer_extra_warp_id)) # Set barrier id for cta sync, epilogue sync and tmem ptr sync self.cta_sync_bar_id = 0 self.epilog_sync_bar_id = 1 @@ -257,6 +308,14 @@ def _setup_attributes(self): ) self.epi_tile_n = cute.size(self.epi_tile[1]) + expected_indexer_epi_n = min(self.mma_tiler[1], 32) + if (self.indexer_q_fusion + and (cute.size(self.epi_tile[0]) != 128 + or self.epi_tile_n != expected_indexer_epi_n)): + raise ValueError( + f"Unexpected swapped indexer-Q epilogue tile {self.epi_tile}; " + f"expected a 128-feature by {expected_indexer_epi_n}-token " + "BF16 stage") # Setup A/B/C stage count in shared memory and ACC stage count in tensor memory self.num_acc_stage, self.num_ab_stage, self.num_c_stage = self._compute_stages( @@ -308,6 +367,10 @@ def _setup_attributes(self): ) self.overlapping_accum = self.num_acc_stage == 1 + if self.indexer_q_fusion and self.overlapping_accum: + raise ValueError( + "The swapped indexer-Q transform-warp schedule requires at " + "least two accumulator stages") sf_atom_mn = 32 self.num_sfa_tmem_cols = (self.cta_tile_shape_mnk[0] // sf_atom_mn) * mma_inst_tile_k @@ -336,6 +399,10 @@ def __call__( max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, + packed_tensor: Optional[cute.Tensor] = None, + indexer_scale_tensor: Optional[cute.Tensor] = None, + position_ids_tensor: Optional[cute.Tensor] = None, + cos_sin_cache_tensor: Optional[cute.Tensor] = None, ): """Execute the GEMM operation in steps: - Setup static attributes before smem/grid/tma computation @@ -368,6 +435,17 @@ def __call__( b_tensor).mma_major_mode() self.c_layout = utils.LayoutEnum.from_tensor(c_tensor) + if cutlass.const_expr( + self.indexer_q_fusion and + (self.c_dtype != cutlass.BFloat16 + or self.c_layout != utils.LayoutEnum.COL_MAJOR + or packed_tensor is None or indexer_scale_tensor is None + or position_ids_tensor is None or cos_sin_cache_tensor is None)): + raise ValueError( + "The swapped indexer-Q epilogue requires a column-major BF16 " + "GEMM boundary plus packed output, scale, position, and RoPE tensors" + ) + # Check if input data types are compatible with MMA instruction if cutlass.const_expr(self.a_dtype != self.b_dtype): raise TypeError( @@ -573,6 +651,10 @@ class SharedStorage: tma_tensor_sfb, tma_atom_c, tma_tensor_c, + packed_tensor, + indexer_scale_tensor, + position_ids_tensor, + cos_sin_cache_tensor, self.cluster_layout_vmnk, self.cluster_layout_sfb_vmnk, self.a_smem_layout_staged, @@ -611,6 +693,10 @@ def kernel( mSFB_nkl: cute.Tensor, tma_atom_c: Optional[cute.CopyAtom], mC_mnl: cute.Tensor, + mPacked_nml: Optional[cute.Tensor], + mIndexerScale_nml: Optional[cute.Tensor], + mPositionIds: Optional[cute.Tensor], + mCosSinCache: Optional[cute.Tensor], cluster_layout_vmnk: cute.Layout, cluster_layout_sfb_vmnk: cute.Layout, a_smem_layout_staged: cute.ComposedLayout, @@ -1430,7 +1516,9 @@ def kernel( "async.shared", space="cta", ) - epilog_threads = 32 * len(self.epilog_warp_id) + epilog_threads = 32 * (self.indexer_transform_warps + if self.indexer_q_fusion else len( + self.epilog_warp_id)) cute.arch.barrier( barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads, @@ -1439,15 +1527,31 @@ def kernel( # # TMA store C to global memory # - if warp_idx == self.epilog_warp_id[0]: - cute.copy( - tma_atom_c, - bSG_sC[(None, c_buffer)], - bSG_gC[(None, real_subtile_idx)], + if cutlass.const_expr(self.indexer_q_fusion): + self._indexer_q_transform_rows( + sC, + c_buffer, + warp_idx, + self.indexer_transform_warps, + cur_tile_coord[0], + cur_tile_coord[1], + cur_tile_coord[2], + real_subtile_idx, + mPacked_nml, + mIndexerScale_nml, + mPositionIds, + mCosSinCache, ) - # Fence and barrier to make sure shared memory store is visible to TMA store - c_pipeline.producer_commit() - c_pipeline.producer_acquire() + else: + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, real_subtile_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() cute.arch.barrier( barrier_id=self.epilog_sync_bar_id, number_of_threads=epilog_threads, @@ -1487,10 +1591,184 @@ def kernel( # # Wait for C store complete # - c_pipeline.producer_tail() + if cutlass.const_expr(not self.indexer_q_fusion): + c_pipeline.producer_tail() + + # Optional row-transform-only warps. The four canonical epilogue + # warps remain the sole TMEM drain owners; these warps join only the + # two shared-stage barriers and process disjoint complete token rows. + if cutlass.const_expr(self.indexer_q_fusion + and self.indexer_transform_warps > 4): + if (warp_idx >= self.indexer_extra_warp_id[0] + and warp_idx <= self.indexer_extra_warp_id[-1]): + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), + cute.arch.grid_dim()) + work_tile = tile_sched.initial_work_tile_info() + subtile_cnt = (self.cta_tile_shape_mnk[1] // self.epi_tile_n) + transform_threads = 32 * self.indexer_transform_warps + transform_warp_idx = (len(self.epilog_warp_id) + warp_idx - + self.indexer_extra_warp_id[0]) + + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + num_prev_subtiles = (tile_sched.num_tiles_executed * + subtile_cnt) + for subtile_idx in cutlass.range(subtile_cnt): + c_buffer = ((num_prev_subtiles + subtile_idx) % + self.num_c_stage) + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=transform_threads, + ) + self._indexer_q_transform_rows( + sC, + c_buffer, + transform_warp_idx, + self.indexer_transform_warps, + cur_tile_coord[0], + cur_tile_coord[1], + cur_tile_coord[2], + subtile_idx, + mPacked_nml, + mIndexerScale_nml, + mPositionIds, + mCosSinCache, + ) + cute.arch.barrier( + barrier_id=self.epilog_sync_bar_id, + number_of_threads=transform_threads, + ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() griddepcontrol_launch_dependents() + @cute.jit + def _indexer_q_transform_rows( + self, + sC: cute.Tensor, + c_buffer: cutlass.Int32, + row_start: cutlass.Int32, + row_stride: cutlass.Constexpr, + head_idx: cutlass.Int32, + token_tile_idx: cutlass.Int32, + batch_idx: cutlass.Int32, + real_subtile_idx: cutlass.Int32, + mPacked_nml: cute.Tensor, + mIndexerScale_nml: cute.Tensor, + mPositionIds: cute.Tensor, + mCosSinCache: cute.Tensor, + ): + """Transform complete BF16 token rows from the shared epilogue tile.""" + lane_idx = cute.arch.lane_idx() + for row in cutlass.range(row_start, + self.epi_tile_n, + row_stride, + unroll_full=True): + token_idx = (token_tile_idx * self.cta_tile_shape_mnk[1] + + real_subtile_idx * self.epi_tile_n + row) + if token_idx < mPositionIds.shape[0]: + values = cute.make_rmem_tensor((4, ), cutlass.Float32) + for value_idx in cutlass.range_constexpr(4): + feature_idx = lane_idx * 4 + value_idx + values[value_idx] = cutlass.Float32(sC[(feature_idx, row, + c_buffer)]) + + if lane_idx >= 16: + position = mPositionIds[token_idx] + pair_base = (lane_idx * 4 - 64) // 2 + for value_idx in cutlass.range_constexpr(0, 4, 2): + cosine = mCosSinCache[position, + pair_base + value_idx // 2] + sine = mCosSinCache[position, + pair_base + value_idx // 2 + 32] + x = values[value_idx] + y = values[value_idx + 1] + values[value_idx] = (cosine * x - sine * y).to( + cutlass.BFloat16).to(cutlass.Float32) + values[value_idx + 1] = (cosine * y + sine * x).to( + cutlass.BFloat16).to(cutlass.Float32) + + amax = cute.arch.fmax( + cute.arch.fmax(values[0], -values[0]), + cute.arch.fmax(values[1], -values[1]), + ) + amax = cute.arch.fmax( + amax, + cute.arch.fmax( + cute.arch.fmax(values[2], -values[2]), + cute.arch.fmax(values[3], -values[3]), + ), + ) + amax = cute.arch.fmax( + amax, cute.arch.shuffle_sync_bfly(amax, offset=1)) + amax = cute.arch.fmax( + amax, cute.arch.shuffle_sync_bfly(amax, offset=2)) + amax = cute.arch.fmax( + amax, cute.arch.shuffle_sync_bfly(amax, offset=4)) + if amax < cutlass.Float32(1.0e-12): + amax = cutlass.Float32(1.0e-12) + + scale_reg = cute.make_rmem_tensor((1, ), cutlass.Float8E8M0FNU) + scale_reg[0] = (amax * cutlass.Float32(1.0 / 6.0)).to( + cutlass.Float8E8M0FNU) + scale_byte = cute.recast_tensor(scale_reg, cutlass.Uint8)[0] + exponent = cutlass.Uint32(scale_byte) + inverse_bits = cutlass.Uint32(0) + if exponent == cutlass.Uint32(254): + inverse_bits = cutlass.Uint32(0x00400000) + else: + inverse_bits = (cutlass.Uint32(254) - exponent) << 23 + inverse_scale = cutlass.Float32( + llvm.bitcast(cutlass.Float32.mlir_type, + inverse_bits.ir_value())) + for value_idx in cutlass.range_constexpr(4): + values[value_idx] = values[value_idx] * inverse_scale + + packed = _indexer_q_pack_fp4x4( + values[0], + values[1], + values[2], + values[3], + ) + midpoint_adjust = cutlass.Uint16(0) + for value_idx in cutlass.range_constexpr(4): + abs_value = cute.arch.fmax(values[value_idx], + -values[value_idx]) + nibble_adjust = cutlass.Uint16(1 << (value_idx * 4)) + if abs_value == cutlass.Float32(0.75): + midpoint_adjust = midpoint_adjust | nibble_adjust + if abs_value == cutlass.Float32(1.75): + midpoint_adjust = midpoint_adjust | nibble_adjust + if abs_value == cutlass.Float32(3.5): + midpoint_adjust = midpoint_adjust | nibble_adjust + packed = packed - midpoint_adjust + magnitude = packed & cutlass.Uint16(0x7777) + nonzero_sign = ((magnitude | (magnitude << 1) + | (magnitude << 2)) + & cutlass.Uint16(0x4444)) << 1 + packed = cutlass.Uint16(packed & (cutlass.Uint16(0x7777) + | nonzero_sign)) + mPacked_nml[head_idx * 32 + lane_idx, token_idx, + batch_idx] = packed + + exponent0 = cutlass.Uint32( + cute.arch.shuffle_sync(exponent, cutlass.Int32(0))) + exponent1 = cutlass.Uint32( + cute.arch.shuffle_sync(exponent, cutlass.Int32(8))) + exponent2 = cutlass.Uint32( + cute.arch.shuffle_sync(exponent, cutlass.Int32(16))) + exponent3 = cutlass.Uint32( + cute.arch.shuffle_sync(exponent, cutlass.Int32(24))) + if lane_idx == 0: + mIndexerScale_nml[ + head_idx, + token_idx, + batch_idx, + ] = (exponent0 | (exponent1 << 8) | (exponent2 << 16) + | (exponent3 << 24)) + def mainloop_s2t_copy_and_partition( self, sSF: cute.Tensor, @@ -1940,7 +2218,7 @@ def is_valid_tensor_alignment( m: cutlass.Int64, n: cutlass.Int64, k: cutlass.Int64, - l: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 - CUTLASS names the batch mode L. ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric], a_major: str, @@ -1992,7 +2270,7 @@ def can_implement( m: cutlass.Int64, n: cutlass.Int64, k: cutlass.Int64, - l: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 - CUTLASS names the batch mode L. a_major: str, b_major: str, c_major: str, @@ -2050,7 +2328,7 @@ def wrapper( sf_m: cutlass.Int64, sf_n: cutlass.Int64, sf_k: cutlass.Int64, - l: cutlass.Constexpr, + l: cutlass.Constexpr, # noqa: E741 - Preserve the generic wrapper API. a_ptr: cute.Pointer, b_ptr: cute.Pointer, a_sf_ptr: cute.Pointer, @@ -2126,6 +2404,102 @@ def wrapper( self(a_tensor, b_tensor, sfa_tensor, sfb_tensor, c_tensor, alpha_tensor, max_active_clusters, current_stream, epilogue_op) + @cute.jit + def wrapper_indexer_q_swap_ab( + self, + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + sf_m: cutlass.Int64, + sf_n: cutlass.Int64, + sf_k: cutlass.Int64, + cos_sin_rows: cutlass.Int64, + batch_count: cutlass.Constexpr, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + a_sf_ptr: cute.Pointer, + b_sf_ptr: cute.Pointer, + packed_ptr: cute.Pointer, + scale_ptr: cute.Pointer, + position_ids_ptr: cute.Pointer, + cos_sin_ptr: cute.Pointer, + alpha_tensor: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + current_stream: cuda.CUstream, + ): + """Run indexer Q with features on MMA-M and tokens on MMA-N. + + The GEMM boundary is a logical column-major BF16 ``[N, M]`` view. + It is never written to global memory; the custom epilogue drains it to + the existing BF16 shared-memory stage, then writes the production + packed-FP4 and four-UE8M0-per-head outputs directly. + """ + # Swap A/B so the fixed 8192 output features occupy hardware MMA-M and + # the small dynamic token count occupies hardware MMA-N. + weight_tensor = cute.make_tensor( + b_ptr, + layout=cute.make_ordered_layout((n, k, batch_count), + order=(1, 0, 2)), + ) + input_tensor = cute.make_tensor( + a_ptr, + layout=cute.make_ordered_layout((m, k, batch_count), + order=(1, 0, 2)), + ) + weight_sf_tensor = cute.make_tensor( + b_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, sf_n, 4, sf_k, batch_count), + order=(2, 1, 4, 0, 3, 5), + ), + ) + input_sf_tensor = cute.make_tensor( + a_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, sf_m, 4, sf_k, batch_count), + order=(2, 1, 4, 0, 3, 5), + ), + ) + + # Only the shape/layout participate in scheduling and TMEM partitioning. + # The custom epilogue does not issue a TMA store through this tensor. + c_boundary_tensor = cute.make_tensor( + cute.recast_ptr(packed_ptr, dtype=cutlass.BFloat16), + layout=cute.make_ordered_layout((n, m, batch_count), + order=(0, 1, 2)), + ) + packed_tensor = cute.make_tensor( + cute.recast_ptr(packed_ptr, dtype=cutlass.Uint16), + layout=cute.make_ordered_layout((n // 4, m, batch_count), + order=(0, 1, 2)), + ) + scale_tensor = cute.make_tensor( + cute.recast_ptr(scale_ptr, dtype=cutlass.Uint32), + layout=cute.make_ordered_layout((n // 128, m, batch_count), + order=(0, 1, 2)), + ) + position_ids_tensor = cute.make_tensor(position_ids_ptr, + cute.make_layout((m, ))) + cos_sin_cache_tensor = cute.make_tensor( + cos_sin_ptr, + layout=cute.make_ordered_layout((cos_sin_rows, 64), order=(1, 0)), + ) + + self( + weight_tensor, + input_tensor, + weight_sf_tensor, + input_sf_tensor, + c_boundary_tensor, + alpha_tensor, + max_active_clusters, + current_stream, + packed_tensor=packed_tensor, + indexer_scale_tensor=scale_tensor, + position_ids_tensor=position_ids_tensor, + cos_sin_cache_tensor=cos_sin_cache_tensor, + ) + @cute.jit def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py index 0d77f4c57fe4..10992874eb31 100644 --- a/tensorrt_llm/_torch/modules/linear.py +++ b/tensorrt_llm/_torch/modules/linear.py @@ -1269,11 +1269,40 @@ def load_weights_fused_gate_up_linear( def transform_weights(self, module: Linear) -> None: super().transform_weights(module) - if (is_sm_100f() and not (module.use_cute_dsl_blockscaling_mm - or module.disable_deep_gemm)) or \ - get_sm_version() == 120: + use_deep_gemm_layout = ( + is_sm_100f() + and not (module.use_cute_dsl_blockscaling_mm + or module.disable_deep_gemm)) or get_sm_version() == 120 + use_indexer_q_cutedsl_layout = (use_deep_gemm_layout and getattr( + module, "use_indexer_q_cutedsl_fusion", False)) + if use_deep_gemm_layout or use_indexer_q_cutedsl_layout: weight, weight_scale = resmooth_to_fp8_e8m0(module.weight, module.weight_scale) + + if use_indexer_q_cutedsl_layout: + # Native SM100 MXF8 MMA consumes one scale per 32 K values. + # The checkpoint/production quantization contract remains + # 128x128: expand each row block and repeat each K scale four + # times, then materialize CUTLASS/CuTe's 128x4 swizzle once at + # weight-load time. + n = weight.shape[0] + scale_cutedsl = weight_scale.repeat_interleave(128, dim=0)[:n] + scale_cutedsl = scale_cutedsl.repeat_interleave(4, dim=1) + scale_cutedsl = torch.ops.trtllm.block_scale_interleave( + scale_cutedsl.to(torch.float8_e8m0fnu).view(torch.uint8)) + module.register_buffer( + "indexer_q_weight_scale_cutedsl", + scale_cutedsl, + persistent=False, + ) + module.register_buffer( + "indexer_q_alpha_cutedsl", + torch.ones((1, ), dtype=torch.float32, + device=weight.device), + persistent=False, + ) + + if use_deep_gemm_layout: transformed_scale = transform_sf_into_required_layout( weight_scale, mn=weight.shape[0], diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index d5d316518f22..be7d663af7a9 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -66,6 +66,7 @@ l0_dgx_b200: - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_shape_dispatch - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_noncontiguous_split - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_fused_cat_fp4_dsv32_prefill_shape + - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused - unittest/_torch/attention/sparse/test_cpp_custom_ops.py::test_indexer_k_cache_gather_contiguous_fp4 - unittest/_torch/attention/sparse/dsa/test_dsa_fp4_indexer.py - condition: diff --git a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py index 20cfaa24805c..09f16a5e7666 100644 --- a/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py +++ b/tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py @@ -17,6 +17,7 @@ - ``torch.ops.trtllm.indexer_k_cache_gather_op`` - ``torch.ops.trtllm.convert_req_index_to_global`` - ``torch.ops.trtllm.fused_cat_fp4`` +- ``torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell`` """ import pytest @@ -488,7 +489,106 @@ def test_convert_req_index_to_global_block_table_padding(): # =================================================================== -# Test 3: fused_cat_fp4 — bit-exact vs DeepGEMM per_token_cast_to_fp4 +# Test 3: native CuTe DSL Indexer-Q projection + RoPE + FP4 fusion +# =================================================================== + + +@skip_pre_blackwell +@pytest.mark.parametrize("num_tokens", [1, 4, 5, 8, 16, 32, 64, 128]) +def test_cute_dsl_fp8_indexer_q_gemm_rope_fp4_matches_unfused(num_tokens): + """The CuTe DSL fusion must match the production chain bit for bit.""" + from _torch.helpers import per_block_cast_to_fp8_e8m0 + + from tensorrt_llm._torch.autotuner import DistributedTuningStrategy, OptimizationProfile + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import CuteDSLIndexerQBlackwellRunner + from tensorrt_llm.quantization.utils import fp8_utils + + torch.manual_seed(2026) + n_heads = 64 + hidden_size = 1536 + output_size = n_heads * HEAD_DIM + max_position = num_tokens * 2 + + qr = torch.randn(num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(output_size, hidden_size, device="cuda", dtype=torch.bfloat16) + weight_fp8, weight_scale = per_block_cast_to_fp8_e8m0(weight) + + weight_scale_deepgemm = fp8_utils.transform_sf_into_required_layout( + weight_scale, + mn=output_size, + k=hidden_size, + recipe=(1, 128, 128), + is_sfa=False, + ) + weight_scale_cutedsl = weight_scale.repeat_interleave(128, dim=0)[:output_size] + weight_scale_cutedsl = weight_scale_cutedsl.repeat_interleave(4, dim=1) + weight_scale_cutedsl = torch.ops.trtllm.block_scale_interleave( + weight_scale_cutedsl.to(torch.float8_e8m0fnu).view(torch.uint8) + ) + + angles = torch.randn(max_position, HEAD_DIM // 4, device="cuda", dtype=torch.float32) + cos_sin_cache = torch.stack((angles.cos(), angles.sin()), dim=1).contiguous() + position_ids = torch.arange(num_tokens, device="cuda", dtype=torch.int32) * 2 + alpha = torch.ones(1, device="cuda", dtype=torch.float32) + + runner = CuteDSLIndexerQBlackwellRunner(use_tvm_ffi=False) + assert runner.tuning_config.exclude_from_cache + assert runner.tuning_config.distributed_tuning_strategy == DistributedTuningStrategy.INDEPENDENT + tactics = runner.get_valid_tactics( + [ + qr, + weight_fp8, + weight_scale_cutedsl, + position_ids, + cos_sin_cache.view(max_position, HEAD_DIM // 2), + alpha, + ], + OptimizationProfile(), + ) + if num_tokens <= 16: + assert all(tactic in tactics for tactic in runner._small_m_tactics) + else: + assert all(tactic[0] != "swap_ab" for tactic in tactics) + assert ("native", (256, 128), (2, 1), False, 0) in tactics + + packed, scale = torch.ops.trtllm.cute_dsl_fp8_indexer_q_gemm_rope_fp4_blackwell( + qr, + weight_fp8, + weight_scale_cutedsl, + position_ids, + cos_sin_cache.view(max_position, HEAD_DIM // 2), + alpha, + use_tvm_ffi=False, + ) + + q_ref = torch.ops.trtllm.fp8_swap_ab_gemm( + qr, + weight_fp8, + weight_scale_deepgemm, + disable_ue8m0_cast=True, + ).view(num_tokens, n_heads, HEAD_DIM) + torch.ops.trtllm.mla_rope_inplace( + q_ref, + position_ids, + cos_sin_cache, + n_heads, + HEAD_DIM // 2, + HEAD_DIM // 2, + False, + False, + ) + packed_ref, scale_ref = torch.ops.trtllm.fused_cat_fp4( + q_ref[..., : HEAD_DIM // 2], q_ref[..., HEAD_DIM // 2 :] + ) + + assert packed.shape == (num_tokens, output_size // 2) + assert scale.shape == (num_tokens, n_heads) + assert torch.equal(packed.view_as(packed_ref), packed_ref) + assert torch.equal(scale.view_as(scale_ref), scale_ref) + + +# =================================================================== +# Test 4: fused_cat_fp4 — bit-exact vs DeepGEMM per_token_cast_to_fp4 # =================================================================== From c74f21311f94e58a9b78e531d58f7cd615054376 Mon Sep 17 00:00:00 2001 From: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> Date: Tue, 21 Jul 2026 21:28:42 -0700 Subject: [PATCH 2/2] [None][fix] Align indexer-Q fallback validation Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 27 +++++++++++++------ 1 file changed, 19 insertions(+), 8 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 1734397446b6..c40ba5ee904a 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -3651,6 +3651,7 @@ def _fake_single_b( 8192, 16384, ) + _INDEXER_Q_POSITION_IDS_INPUT_INDEX = 3 def _map_cutedsl_indexer_q_tuning_bucket(num_tokens: int) -> int: if num_tokens <= 4: @@ -3663,7 +3664,12 @@ def _map_cutedsl_indexer_q_tuning_bucket(num_tokens: int) -> int: def _prepare_cutedsl_indexer_q_tuning_inputs( inputs: List[torch.Tensor]) -> List[torch.Tensor]: - inputs[3] = torch.zeros_like(inputs[3]) + # The autotuner resizes position_ids to the token bucket, leaving newly + # allocated values uninitialized. Use position zero for every tuning + # row so the fused RoPE lookup always stays inside cos_sin_cache. + position_ids = inputs[_INDEXER_Q_POSITION_IDS_INPUT_INDEX] + inputs[_INDEXER_Q_POSITION_IDS_INPUT_INDEX] = torch.zeros_like( + position_ids) return inputs class CuteDSLIndexerQBlackwellRunner(TunableRunner): @@ -3726,7 +3732,7 @@ def get_valid_tactics( m, k = inputs[0].shape n = inputs[1].shape[0] tactics = [] - if 0 < m <= 16 and n % 128 == 0 and k % 128 == 0: + if self._small_m_kernel_is_supported(m, n, k): tactics.extend(self.__class__._small_m_tactics) tactics.extend([ @@ -3750,12 +3756,17 @@ def get_valid_tactics( return tactics @staticmethod - def _fallback_tactic(m: int) -> Tuple: + def _small_m_kernel_is_supported(m: int, n: int, k: int) -> bool: + return 0 < m <= 16 and n % 128 == 0 and k % 128 == 0 + + @classmethod + def _fallback_tactic(cls, m: int, n: int, k: int) -> Tuple: """Safe eager-mode fallback when TRT-LLM autotuning is disabled.""" - if m <= 4: - return ("swap_ab", (128, 16), (1, 1), False, 4) - if m <= 8: - return ("swap_ab", (128, 16), (1, 1), False, 8) + if cls._small_m_kernel_is_supported(m, n, k): + if m <= 4: + return ("swap_ab", (128, 16), (1, 1), False, 4) + if m <= 8: + return ("swap_ab", (128, 16), (1, 1), False, 8) return ("native", (256, 128), (2, 1), False, 0) @staticmethod @@ -3776,7 +3787,7 @@ def forward( m, k = input.shape n = weight.shape[0] if tactic == -1: - tactic = self._fallback_tactic(m) + tactic = self._fallback_tactic(m, n, k) (kernel_kind, mma_tiler_mn, cluster_shape_mn, use_prefetch, transform_warps) = tactic if kernel_kind == "swap_ab":