feat(kernel): add quanted silu fusion kernel

python/sglang/srt/layers/moe/moe_runner/deep_gemm.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+import os
 from typing import TYPE_CHECKING, List, Optional
 
 import torch
@@ -42,7 +43,18 @@ _is_hip = is_hip()
 _is_npu = is_npu()
 _is_cuda = is_cuda()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
-
+_USE_TAI_FUSED_CONTIG = _is_cuda and get_bool_env_var(
+    "SGLANG_DEEPGEMM_USE_TAI_FUSED_CONTIG"
+)
+_TAI_FUSED_CONTIG_MODE = os.getenv(
+    "SGLANG_DEEPGEMM_TAI_FUSED_CONTIG_MODE", "two_step"
+).lower()
+_USE_TAI_FUSED_CONTIG_THREE_STEP = (
+    _USE_TAI_FUSED_CONTIG and _TAI_FUSED_CONTIG_MODE == "three_step"
+)
+print(f"SGLANG_DEEPGEMM_USE_TAI_FUSED_CONTIG: {_USE_TAI_FUSED_CONTIG}, "
+    f"_USE_TAI_FUSED_CONTIG_THREE_STEP: {_USE_TAI_FUSED_CONTIG_THREE_STEP}"
+)
 if not (_is_npu or _is_hip) and _is_cuda:
     from sgl_kernel import silu_and_mul
 
@@ -76,6 +88,45 @@ def copy_list_to_gpu_no_ce(arr: List[int]):
     return tensor_gpu
 
 
+if _USE_TAI_FUSED_CONTIG:
+    from tai_kernel.quantization import (
+        fused_silu_and_mul_quant_fp8 as _contig_post_quant_fp8,
+    )
+
+    if _USE_TAI_FUSED_CONTIG_THREE_STEP:
+        from tai_kernel.quantization import (
+            fused_silu_and_mul_quant_fp8_three_step as _contig_post_quant_fp8,
+        )
+else:
+
+    def _contig_post_quant_fp8(
+        x: torch.Tensor,
+        group_size: int,
+        column_major_scales: bool,
+        scale_tma_aligned: bool,
+        scale_ue8m0: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        from sglang.srt.layers.quantization.fp8_kernel import (
+            sglang_per_token_group_quant_fp8,
+        )
+
+        down_input = torch.empty(
+            x.shape[:-1] + (x.shape[-1] // 2,),
+            device=x.device,
+            dtype=torch.bfloat16,
+        )
+        silu_and_mul(x.view(-1, x.shape[-1]), down_input.view(-1, down_input.shape[-1]))
+        x_q, x_s = sglang_per_token_group_quant_fp8(
+            x=down_input,
+            group_size=group_size,
+            column_major_scales=column_major_scales,
+            scale_tma_aligned=scale_tma_aligned,
+            scale_ue8m0=scale_ue8m0,
+        )
+        del down_input
+        return x_q, x_s
+
+
 @dataclass
 class DeepGemmRunnerInput(RunnerInput):
     hidden_states: torch.Tensor
@@ -138,9 +189,6 @@ class DeepGemmRunnerCore(MoeRunnerCore):
         running_state: dict,
     ) -> torch.Tensor:
         from sglang.srt.layers.moe.ep_moe.kernels import tma_align_input_scale
-        from sglang.srt.layers.quantization.fp8_kernel import (
-            sglang_per_token_group_quant_fp8,
-        )
 
         hidden_states = runner_input.hidden_states
         hidden_states_scale = runner_input.hidden_states_scale
@@ -176,33 +224,23 @@ class DeepGemmRunnerCore(MoeRunnerCore):
 
         dispose_tensor(hidden_states)
         dispose_tensor(hidden_states_scale)
-
-        down_input = torch.empty(
-            (
-                all_tokens,
-                N // 2,
-            ),
-            device=gateup_output.device,
-            dtype=torch.bfloat16,
-        )
-        silu_and_mul(gateup_output.view(-1, N), down_input)
-        del gateup_output
-
-        down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8(
-            down_input,
-            scale_block_size,
+        down_input_fp8, down_input_scale = _contig_post_quant_fp8(
+            x=gateup_output,
+            group_size=scale_block_size,
             column_major_scales=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
             scale_tma_aligned=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
             scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
         )
-        del down_input
-
+        del gateup_output
         down_output = torch.empty(
             (all_tokens, K),
             device=hidden_states_device,
             dtype=torch.bfloat16,
         )
-        if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+        if (
+            not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
+            and not _USE_TAI_FUSED_CONTIG_THREE_STEP
+        ):
             down_input_scale = tma_align_input_scale(down_input_scale)
 
         deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_contig(