[Diffusion] Fix weight scale swizzle and add large-M kernel config for FLUX.2-dev-NVFP4 (#22064)

python/sglang/jit_kernel/csrc/gemm/nvfp4/nvfp4_scaled_mm_sm100.cuh

@@ -53,7 +53,7 @@ const dim3 KernelConfigM256<T>::preferred_cluster(2, 4, 1);
 template <typename T>
 const dim3 KernelConfigM256<T>::fallback_cluster(2, 1, 1);
 
-// Default config(half_t/bfloat16_t) for M > 256
+// Config(half_t/bfloat16_t) for 256 < M <= 1024
 template <typename T>
 struct KernelConfigDefault {
   using OutputType = T;
@@ -66,10 +66,27 @@ struct KernelConfigDefault {
   const static dim3 fallback_cluster;
 };
 template <typename T>
-const dim3 KernelConfigDefault<T>::preferred_cluster(4, 4, 1);
+const dim3 KernelConfigDefault<T>::preferred_cluster(2, 4, 1);
 template <typename T>
 const dim3 KernelConfigDefault<T>::fallback_cluster(2, 1, 1);
 
+// Config(half_t/bfloat16_t) for M > 1024: 1x4 cluster reduces M-tail waste.
+template <typename T>
+struct KernelConfigLargeM {
+  using OutputType = T;
+  using MmaTileShape = Shape<_256, _256, _256>;
+  using ClusterShape = Shape<int, int, _1>;
+  using EpilogueTile = Shape<_128, _64>;
+  using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized2Sm;
+  using MainloopSchedule = cutlass::gemm::KernelTmaWarpSpecialized2SmNvf4Sm100;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+template <typename T>
+const dim3 KernelConfigLargeM<T>::preferred_cluster(1, 4, 1);
+template <typename T>
+const dim3 KernelConfigLargeM<T>::fallback_cluster(1, 2, 1);
+
 struct KernelConfigFp32 {
   using OutputType = float;
   using MmaTileShape = Shape<_128, _128, _256>;
@@ -261,8 +278,12 @@ void cutlassFp4GemmDispatchSm100(
     runGemm<Fp4GemmSm100<KernelConfigM128<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
   } else if (m <= 256) {
     runGemm<Fp4GemmSm100<KernelConfigM256<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
-  } else {
+  } else if (m <= 1024) {
+    // m in (256, 1024]: 2x4 cluster balances SM occupancy and data reuse
     runGemm<Fp4GemmSm100<KernelConfigDefault<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
+  } else {
+    // m in (1024, inf): 1x4 cluster eliminates M-tail waste for FLUX-class shapes
+    runGemm<Fp4GemmSm100<KernelConfigLargeM<OutType>>>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream);
   }
 }
 

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py

@@ -43,16 +43,6 @@ def _get_fp4_gemm_op():
     return current_platform.get_modelopt_fp4_gemm_op()
 
 
-@lru_cache(maxsize=1)
-def _get_comfy_kitchen_cuda_backend():
-    try:
-        import comfy_kitchen.backends.cuda as ck_cuda
-
-        return ck_cuda
-    except Exception:
-        return None
-
-
 class ModelOptQuantConfig(QuantizationConfig):
     def __init__(
         self,
@@ -224,20 +214,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
         return False
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
-        should_use_best_perf_kit = getattr(
+        return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod)
-            current_platform, "should_use_modelopt_fp4_best_performance_kit", None
-        )
-        warn_missing_best_perf_kit = getattr(
-            current_platform, "warn_if_modelopt_fp4_best_performance_kit_missing", None
-        )
-
-        if callable(should_use_best_perf_kit) and should_use_best_perf_kit():
-            linear_cls = ComfyUIFp4LinearMethod
-        else:
-            if callable(warn_missing_best_perf_kit):
-                warn_missing_best_perf_kit()
-            linear_cls = ModelOptFp4LinearMethod
-        return self._get_quant_method(layer, prefix, Linear=linear_cls)
 
 
 class ModelOptFp4LinearMethod(LinearMethodBase):
@@ -349,6 +326,11 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         K_padded = round_up(K, 4)
         padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype)
         padded_scales[:B, :M, :K] = scales
+        # Blockwise interleave for CUTLASS TMA layout required by CUTLASS kernel
+        padded_scales = padded_scales.reshape(
+            B, M_padded // 128, 4, 32, K_padded // 4, 4
+        )
+        padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5)
         padded_scales = padded_scales.contiguous().cuda()
         padded_scales = (
             padded_scales.reshape(M_padded, K_padded)
@@ -417,146 +399,3 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         if bias is not None:
             out = out + bias
         return out.view(*output_shape)
-
-
-class ComfyUIFp4LinearMethod(LinearMethodBase):
-    """NVFP4 linear method using comfy-kitchen cuBLAS kernels (Blackwell)."""
-
-    def __init__(self, quant_config: ModelOptFp4Config):
-        self.quant_config = quant_config
-
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        input_size_per_partition: int,
-        output_partition_sizes: List[int],
-        input_size: int,
-        output_size: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        del input_size, output_size
-        if not self.quant_config.is_checkpoint_nvfp4_serialized:
-            raise ValueError(
-                "NVFP4 quantization was selected, "
-                "dynamic quantization is not supported."
-            )
-        if input_size_per_partition % 16 != 0:
-            raise ValueError(
-                f"Unsupported model when input features size is {input_size_per_partition}, "
-                "not multiple of 16, for NVFP4 quantization."
-            )
-
-        output_size_per_partition = sum(output_partition_sizes)
-        weight_loader = extra_weight_attrs.get("weight_loader")
-
-        layer.logical_widths = output_partition_sizes
-        layer.input_size_per_partition = input_size_per_partition
-        layer.output_size_per_partition = output_size_per_partition
-
-        weight = ModelWeightParameter(
-            data=torch.empty(
-                output_size_per_partition,
-                input_size_per_partition // 2,
-                dtype=torch.uint8,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight", weight)
-
-        input_scale = PerTensorScaleParameter(
-            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
-            weight_loader=weight_loader,
-        )
-        set_weight_attrs(input_scale, {"missing_param_init": "ones"})
-        layer.register_parameter("input_scale", input_scale)
-
-        weight_scale_2 = PerTensorScaleParameter(
-            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight_scale_2", weight_scale_2)
-
-        weight_scale = ModelWeightParameter(
-            data=torch.empty(
-                output_size_per_partition,
-                input_size_per_partition // self.quant_config.group_size,
-                dtype=torch.float8_e4m3fn,
-            ),
-            input_dim=1,
-            output_dim=0,
-            weight_loader=weight_loader,
-        )
-        layer.register_parameter("weight_scale", weight_scale)
-
-    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-        from comfy_kitchen.float_utils import from_blocked, to_blocked
-
-        input_scale = layer.input_scale.max().to(torch.float32)
-        weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
-
-        copy_or_rebind_param(layer, "input_scale_ck", input_scale.cuda())
-        copy_or_rebind_param(layer, "weight_scale_2_ck", weight_scale_2.cuda())
-        layer.output_size_per_partition = layer.weight.shape[0]
-        copy_or_rebind_param(layer, "weight", layer.weight.data.contiguous().cuda())
-
-        # Checkpoint block scales are already in cuBLAS tiled layout.
-        # Pad to (roundup(N, 128), roundup(K//16, 4)) if needed.
-        scales = layer.weight_scale.data
-        N, Ks = scales.shape
-        N_padded = round_up(N, 128)
-        Ks_padded = round_up(Ks, 4)
-
-        if N == N_padded and Ks == Ks_padded:
-            weight_scale_ck = scales.cuda()
-        else:
-            scales_rm = from_blocked(scales, num_rows=N, num_cols=Ks)
-            padded_rm = torch.zeros((N_padded, Ks_padded), dtype=scales.dtype)
-            padded_rm[:N, :Ks] = scales_rm
-            weight_scale_ck = to_blocked(padded_rm, flatten=False).cuda()
-
-        copy_or_rebind_param(layer, "weight_scale_ck", weight_scale_ck)
-
-    def apply(
-        self,
-        layer: torch.nn.Module,
-        x: torch.Tensor,
-        bias: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        ck_cuda = _get_comfy_kitchen_cuda_backend()
-        if ck_cuda is None:
-            raise RuntimeError(
-                "comfy_kitchen is not available. "
-                "Install it to use ComfyUIFp4LinearMethod."
-            )
-
-        output_dtype = x.dtype
-        input_shape = x.shape
-        x_2d = x.view(-1, input_shape[-1])  # [M, K]
-        M = x_2d.shape[0]
-
-        output_size = layer.output_size_per_partition
-        output_shape = list(input_shape[:-1]) + [output_size]
-
-        if not x_2d.is_contiguous():
-            x_2d = x_2d.contiguous()
-
-        x_fp4, x_block_scale = ck_cuda.quantize_nvfp4(
-            x_2d, layer.input_scale_ck, pad_16x=True
-        )
-
-        out = ck_cuda.scaled_mm_nvfp4(
-            x_fp4,
-            layer.weight,
-            tensor_scale_a=layer.input_scale_ck,
-            tensor_scale_b=layer.weight_scale_2_ck,
-            block_scale_a=x_block_scale,
-            block_scale_b=layer.weight_scale_ck,
-            bias=bias,
-            out_dtype=output_dtype,
-        )
-        out = out[:M, :output_size].contiguous()
-
-        return out.view(*output_shape)

python/sglang/multimodal_gen/runtime/platforms/cuda.py

@@ -145,53 +145,6 @@ class CudaPlatformBase(Platform):
         except ImportError:
             return None, None
 
-    @classmethod
-    @lru_cache(maxsize=1)
-    def has_modelopt_fp4_best_performance_kit(cls) -> bool:
-        try:
-            import comfy_kitchen.backends.cuda  # noqa: F401
-
-            return True
-        except Exception:
-            return False
-
-    @classmethod
-    @lru_cache(maxsize=1)
-    def can_use_modelopt_fp4_best_performance_kit(cls) -> bool:
-        if not cls.is_blackwell() or not cls.has_modelopt_fp4_best_performance_kit():
-            return False
-
-        try:
-            import comfy_kitchen.backends.cuda as ck_cuda
-
-            device = cls.get_local_torch_device()
-            x = torch.zeros((16, 16), dtype=torch.bfloat16, device=device)
-            scale = torch.ones((), dtype=torch.float32, device=device)
-            ck_cuda.quantize_nvfp4(x, scale, pad_16x=True)
-            return True
-        except Exception as e:
-            logger.warning(
-                "best performance kit (comfy-kitchen) is installed but unusable on "
-                "this system (%s). Blackwell NVFP4 will fall back to the generic "
-                "ModelOpt FP4 path.",
-                e,
-            )
-            return False
-
-    @classmethod
-    def should_use_modelopt_fp4_best_performance_kit(cls) -> bool:
-        return cls.can_use_modelopt_fp4_best_performance_kit()
-
-    @classmethod
-    @lru_cache(maxsize=1)
-    def warn_if_modelopt_fp4_best_performance_kit_missing(cls) -> None:
-        if cls.is_blackwell() and not cls.has_modelopt_fp4_best_performance_kit():
-            logger.warning(
-                "best performance kit (comfy-kitchen) is not installed. "
-                "Blackwell NVFP4 will fall back to the generic ModelOpt FP4 path. "
-                "Install it with `pip install comfy-kitchen[cublas]`."
-            )
-
     @classmethod
     def is_full_nvlink(cls, device_ids: list[int]) -> bool:
         raise NotImplementedError