Fix rebase issues

Signed-off-by: Hao Lu <[email protected]>

Author: hlu1

cpp/tensorrt_llm/kernels/trtllmGenKernels/CMakeLists.txt

@@ -17,4 +17,4 @@
 
 add_subdirectory(fmha)
 add_subdirectory(blockscaleGemm)
-add_subdirectory(fp8BlockScaleMoe)
+add_subdirectory(blockScaleMoe)

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/CMakeLists.txt renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/CMakeLists.txt

@@ -15,7 +15,7 @@
  */
 
 #include "tensorrt_llm/kernels/quantization.h"
-#include "tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/runner.h"
+#include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h"
 #include "tensorrt_llm/runtime/torchUtils.h"
 #include "tensorrt_llm/thop/thUtils.h"
 #include <ATen/cuda/EmptyTensor.h>

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/gemmCommon.h renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/gemmCommon.h

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-#include "tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/runner.h"
+#include "tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.h"
 #include "tensorrt_llm/runtime/torchUtils.h"
 #include "tensorrt_llm/thop/thUtils.h"
 #include <ATen/cuda/EmptyTensor.h>

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/gemmCubins/MoE_ProjDown_BatchN_E2m1Fp32_Bfloat16_Tile128x8x512_EpiTile128x8_Mma128x8x64_Cluster1x1x1_transposeMmaOutput_sm100a_cubin.h renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/gemmCubins/MoE_ProjDown_BatchN_E2m1Fp32_Bfloat16_Tile128x8x512_EpiTile128x8_Mma128x8x64_Cluster1x1x1_transposeMmaOutput_sm100a_cubin.h

@@ -241,6 +241,10 @@ def parse_args():
         default='TRTLLM',
         choices=['TRTLLM', 'FLASHINFER'],
         help='Attention kernel for PyTorch. Ignored for TRT backend.')
+    parser.add_argument('--moe_backend',
+                        type=str,
+                        default='CUTLASS',
+                        choices=['CUTLASS', 'TRTLLM'])
     parser.add_argument("--enable_chunked_prefill",
                         action="store_true",
                         help="Exercises the chunked prefill inference feature.")
@@ -307,7 +311,7 @@ def main():
         assert args.engine_dir is None, "pytorch backend does not need TRT Engine"
         config = PyTorchConfig(
             attn_backend=args.attn_backend,
-            moe_backend='TRTLLM',
+            moe_backend=args.moe_backend,
             enable_overlap_scheduler=args.enable_overlap_scheduler,
             torch_compile_enabled=args.torch_compile)
         llm = tensorrt_llm._torch.LLM(

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/gemmCubins/MoE_ProjDown_BatchN_E4m3Fp32_Bfloat16_Tile128x8x128_EpiTile64x8_Mma64x8x32_Cluster1x1x1_transposeMmaOutput_DsFp8_sm100a_cubin.h renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/gemmCubins/MoE_ProjDown_BatchN_E4m3Fp32_Bfloat16_Tile128x8x128_EpiTile64x8_Mma64x8x32_Cluster1x1x1_transposeMmaOutput_DsFp8_sm100a_cubin.h

@@ -361,6 +361,7 @@ def __init__(
 
         if quant_mode.has_fp8_block_scales():
             mla_weight_dtype = torch.float8_e4m3fn
+            # TODO: remove hack for fp8 Deepseek on SM100
             if config.moe_backend == "TRTLLM":
                 mla_weight_dtype = dtype
         else:
@@ -487,7 +488,6 @@ def forward(
         attn_metadata: AttentionMetadata,
         all_reduce_params: Optional[AllReduceParams] = None,
     ) -> torch.Tensor:
-        assert hidden_states.dtype == torch.bfloat16, "Just for TRTLLM FP8 E2E test"
         if self.is_lite:
             compressed_kv, k_pe = self.fused_a(hidden_states).split(
                 [self.kv_lora_rank, self.qk_rope_head_dim], -1)

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/gemmCubins/MoE_ProjUp_BatchN_E2m1Fp32_E2m1_Tile128x8x512_EpiTile128x8_Mma128x8x64_Cluster1x1x1_transposeMmaOutput_InplaceRoute_GatedAct_sm100a_cubin.h renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/gemmCubins/MoE_ProjUp_BatchN_E2m1Fp32_E2m1_Tile128x8x512_EpiTile128x8_Mma128x8x64_Cluster1x1x1_transposeMmaOutput_InplaceRoute_GatedAct_sm100a_cubin.h

@@ -335,6 +335,7 @@ def setup_quant_scales(self):
                 fc2_weight_block=self.w2_weight_scale,
                 fc2_global=self.fc2_alpha,
             )
+
     def is_trtllm(self):
         return self.moe_backend == "TRTLLM" and self.quant_config is not None
 
@@ -416,6 +417,7 @@ def create_weights(self):
                 self.register_parameter("w2_weight_scaling_factor",
                                         w2_weight_scaling_factor)
             elif qc.quant_mode.has_nvfp4():
+                self.has_nv_fp4 = True
                 if self.is_trtllm():
                     weight_dtype = float4_sf_dtype
                     weight_vec_size = torch.iinfo(weight_dtype).bits // 4
@@ -668,7 +670,8 @@ def forward(
         all_rank_num_tokens: Optional[List[int]] = None,
     ) -> torch.Tensor:
         if self.is_cutlass():
-            return self.forward_cutlass(x, router_logits, min_latency_mode, output_dtype, all_rank_num_tokens)
+            return self.forward_cutlass(x, router_logits, min_latency_mode,
+                                        output_dtype, all_rank_num_tokens)
         elif self.is_trtllm():
             return self.forward_trtllmgen(x, router_logits)
         else:
@@ -763,14 +766,7 @@ def forward_trtllmgen(self, x: torch.Tensor,
 
         if self.quant_config and self.quant_config.quant_mode.has_fp8_block_scales(
         ):
-            # TODO: We need a new kernel to support fp8 block scaling for blackwell
             x_val, x_scale = torch.ops.trtllm.fp8_quantize_1x128(x)
-            m_4_align = (x.shape[0] + 3) // 4 * 4
-            kscal_128 = (x.shape[1] + 127) // 128
-            act_scal_elesize = kscal_128 * m_4_align
-            x_scale = x_scale[:act_scal_elesize]
-            x_scale = x_scale.view(kscal_128, m_4_align)
-            x_scale = x_scale[:kscal_128, :x.shape[0]].contiguous()
 
             final_hidden_states = torch.ops.trtllm.fp8_block_scale_moe_runner(
                 router_logits,

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/gemmCubins/MoE_ProjUp_BatchN_E4m3Fp32_E4m3_Tile128x8x128_EpiTile64x8_Mma64x8x32_Cluster1x1x1_transposeMmaOutput_DsFp8_InplaceRoute_sm100a_cubin.h renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/gemmCubins/MoE_ProjUp_BatchN_E4m3Fp32_E4m3_Tile128x8x128_EpiTile64x8_Mma64x8x32_Cluster1x1x1_transposeMmaOutput_DsFp8_InplaceRoute_sm100a_cubin.h

@@ -322,15 +322,9 @@ def apply_linear(self, input, weight, bias):
                 if input.dtype == torch.float8_e4m3fn:
                     input = input.to(torch.bfloat16) * self.input_scale
                 assert input.dtype == torch.bfloat16
-                # TODO: We need a new kernel to support fp8 block scaling for blackwell
-                act_input_fp8, a_scale = torch.ops.trtllm.fp8_quantize_1x128(
+
+                act_input_fp8, act_input_sf = torch.ops.trtllm.fp8_quantize_1x128(
                     input)
-                m_4_align = (input.shape[0] + 3) // 4 * 4
-                kscal_128 = (input.shape[1] + 127) // 128
-                act_scal_elesize = kscal_128 * m_4_align
-                a_scale = a_scale[:act_scal_elesize]
-                a_scale = a_scale.view(kscal_128, m_4_align)
-                act_input_sf = a_scale[:kscal_128, :input.shape[0]].contiguous()
 
                 output = torch.ops.trtllm.fp8_block_scaling_gemm(
                     act_input_fp8, self.weight, act_input_sf, self.weight_scale)
@@ -402,11 +396,7 @@ def load_weights(self, weights: List[Dict]):
         assert self._weights_created
 
         def copy(dst: Parameter, src: torch.Tensor):
-            # TODO: Update this once we have BMM FP8 working with blackwell
-            #assert dst.dtype == src.dtype, f"Incompatible dtype. dst: {dst.dtype}, src: {src.dtype}"
-            assert dst.dtype == src.dtype or (
-                dst.dtype == torch.bfloat16 and src.dtype == torch.float8_e4m3fn
-            ), f"Incompatible dtype. dst: {dst.dtype}, src: {src.dtype}"
+            assert dst.dtype == src.dtype, f"Incompatible dtype. dst: {dst.dtype}, src: {src.dtype}"
             dst.data.copy_(src)
 
         weight_mode = self.weights_loading_config.weight_mode

cpp/tensorrt_llm/kernels/trtllmGenKernels/fp8BlockScaleMoe/gemmList.h renamed to cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/gemmList.h

@@ -183,7 +183,7 @@ def create_py_executor(executor_config: ExecutorConfig,
         executor_config.pytorch_backend_config.use_kv_cache = False
 
     kv_cache_max_tokens = None
-    # TODO: remove this once we have a loop fix for routing token limit
+
     if model_engine.model.model_config.is_generation:
         kv_cache_max_tokens = estimate_max_kv_cache_tokens(
             model_engine, executor_config, mapping)