[None][perf] Fuse DeepSeek-V4 Indexer Q projection with CuTe DSL

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 <int WarpsPerBlock>
+template <int WarpsPerBlock, bool OutputCuteDslSf>
 __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<uint32_t>(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<uint32_t>(ue8m0_scale.__x), 0);
-        uint32_t const s1 = __shfl_sync(0xFFFFFFFFu, static_cast<uint32_t>(ue8m0_scale.__x), 8);
-        uint32_t const s2 = __shfl_sync(0xFFFFFFFFu, static_cast<uint32_t>(ue8m0_scale.__x), 16);
-        uint32_t const s3 = __shfl_sync(0xFFFFFFFFu, static_cast<uint32_t>(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<int64_t>(m_idx / 128) * num_k_tiles * 512
+                    + static_cast<int64_t>(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<uint32_t*>(reinterpret_cast<uint8_t*>(packed_scale_output) + dst_offset) = replicated;
+            }
+        }
+    }
+    else
     {
-        packed_scale_output[static_cast<int64_t>(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<int64_t>(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<kWarpsPerBlock>, grid, block, 0, stream, fp8_output, packed_scale_output,
-        input, m, k, scale_leading_dim_uint32);
+        fp8_quantize_1x128_packed_kernel_impl<kWarpsPerBlock, false>, 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<kWarpsPerBlock, true>, grid, block, 0, stream, fp8_output,
+        reinterpret_cast<int32_t*>(swizzled_scale_output), input, m, k, padded_m);
 }
 
 } // namespace kernels::fp8_blockscale_gemm

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

cpp/tensorrt_llm/thop/fp8Quantize.cpp

@@ -209,6 +209,38 @@ std::tuple<at::Tensor, at::Tensor> fp8_quantize_1x128_packed_ue8m0(at::Tensor co
 
     return {valueE4M3.slice(0, 0, m), packedScale};
 }
+
+std::tuple<at::Tensor, at::Tensor> 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<int32_t>::max(), "M must be within int32");
+    TORCH_CHECK(k <= std::numeric_limits<int32_t>::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<uint8_t>(),
+        reinterpret_cast<__nv_bfloat16 const*>(self.data_ptr()), static_cast<int>(m), static_cast<int>(k),
+        static_cast<int>(paddedM), stream);
+
+    return {valueE4M3, scaleE8M0};
+}
 } // namespace torch_ext
 
 TRTLLM_NAMESPACE_END
@@ -218,11 +250,13 @@ 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)
 {
     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);
 }

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
@@ -1079,6 +1079,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,
@@ -1138,6 +1143,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(
@@ -1152,6 +1161,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
@@ -1281,7 +1333,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()
@@ -1302,9 +1354,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)
@@ -1321,9 +1371,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)
 

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,