chore: import upstream snapshot with attribution

python/sglang/jit_kernel/trtllm_lora_temp/SOURCE.md

@@ -0,0 +1,20 @@
+FlashInfer TRTLLM MoE Overlay
+=============================
+
+This directory contains only the editable overlay files used by SGLang's
+`experimental_sgl_trtllm` SM100 TRTLLM fused MoE backend. Unmodified FlashInfer
+and TRTLLM sources are compiled from the installed `flashinfer` package at JIT
+time.
+
+Local overlay source:
+
+- `data/csrc/trtllm_fused_moe_kernel_launcher.cu`
+- `data/csrc/trtllm_fused_moe_runner.cu`
+- `data/csrc/fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu`
+- `data/include/flashinfer/trtllm/fused_moe/DevKernel.h`
+- `data/include/flashinfer/trtllm/fused_moe/runner.h`
+
+The backend still depends on the installed `flashinfer` and `flashinfer_cubin`
+packages for the rest of the FlashInfer/TRTLLM JIT source tree and TRTLLM-Gen
+BMM cubin artifacts. The local include directory is passed before FlashInfer's
+installed include directory so these overlay headers shadow the originals.

python/sglang/jit_kernel/trtllm_lora_temp/__init__.py

@@ -0,0 +1,19 @@
+from sglang.jit_kernel.trtllm_lora_temp.core import (
+    trtllm_bf16_routed_moe_lora,
+    trtllm_fp4_block_scale_moe_lora_finalize,
+    trtllm_fp4_block_scale_routed_moe_lora,
+    trtllm_fp8_block_scale_moe,
+    trtllm_fp8_block_scale_moe_lora_finalize,
+    trtllm_fp8_block_scale_routed_moe,
+    trtllm_fp8_block_scale_routed_moe_lora,
+)
+
+__all__ = [
+    "trtllm_bf16_routed_moe_lora",
+    "trtllm_fp4_block_scale_moe_lora_finalize",
+    "trtllm_fp4_block_scale_routed_moe_lora",
+    "trtllm_fp8_block_scale_moe_lora_finalize",
+    "trtllm_fp8_block_scale_moe",
+    "trtllm_fp8_block_scale_routed_moe",
+    "trtllm_fp8_block_scale_routed_moe_lora",
+]

python/sglang/jit_kernel/trtllm_lora_temp/core.py

@@ -0,0 +1,514 @@
+import functools
+from typing import List, Optional, Union
+
+import torch
+
+from sglang.srt.lora.trtllm_lora_temp.environ import lora_envs
+
+
+@functools.cache
+def get_sgl_trtllm_moe_sm100_module():
+    import flashinfer.fused_moe.core as fi_core
+
+    from sglang.jit_kernel.trtllm_lora_temp.jit import (
+        gen_sgl_trtllm_gen_fused_moe_sm100_module,
+    )
+
+    original_gen = fi_core.gen_trtllm_gen_fused_moe_sm100_module
+    fi_core.gen_trtllm_gen_fused_moe_sm100_module = (
+        gen_sgl_trtllm_gen_fused_moe_sm100_module
+    )
+    try:
+        fi_core.get_trtllm_moe_sm100_module.cache_clear()
+        return fi_core.get_trtllm_moe_sm100_module()
+    finally:
+        fi_core.gen_trtllm_gen_fused_moe_sm100_module = original_gen
+        fi_core.get_trtllm_moe_sm100_module.cache_clear()
+
+
+@functools.cache
+def get_sgl_trtllm_moe_sm100_raw_module():
+    from flashinfer.fused_moe.core import setup_cubin_loader
+
+    from sglang.jit_kernel.trtllm_lora_temp.jit import (
+        gen_sgl_trtllm_gen_fused_moe_sm100_module,
+    )
+
+    module = gen_sgl_trtllm_gen_fused_moe_sm100_module()
+    moe_op = module.build_and_load()
+    setup_cubin_loader(str(module.get_library_path()))
+    return moe_op
+
+
+def _validate_routing_replay_out(
+    routing_replay_out: Optional[torch.Tensor], top_k: int
+) -> None:
+    if routing_replay_out is None:
+        return
+    assert routing_replay_out.dim() == 2
+    assert routing_replay_out.shape[1] == top_k
+    assert routing_replay_out.dtype == torch.int16
+    assert routing_replay_out.is_cuda
+    assert routing_replay_out.is_contiguous()
+
+
+def trtllm_fp8_block_scale_moe(
+    routing_logits: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    do_finalize: bool = True,
+    enable_pdl: Optional[bool] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type=None,
+    activation_type: Optional[int] = None,
+    norm_topk_prob: bool = True,
+    routing_replay_out: Optional[torch.Tensor] = None,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    from flashinfer.fused_moe.core import ActivationType, Fp8QuantizationType
+
+    _validate_routing_replay_out(routing_replay_out, top_k)
+
+    if fp8_quantization_type is None:
+        fp8_quantization_type = Fp8QuantizationType.DeepSeekFp8
+    if activation_type is None:
+        activation_type = ActivationType.Swiglu.value
+
+    output = torch.empty(
+        hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
+    )
+    result = get_sgl_trtllm_moe_sm100_module().trtllm_fp8_block_scale_moe(
+        routing_logits,
+        None,
+        None,
+        routing_bias,
+        hidden_states,
+        hidden_states_scale,
+        gemm1_weights,
+        gemm1_weights_scale,
+        gemm2_weights,
+        gemm2_weights_scale,
+        output,
+        num_experts,
+        top_k,
+        n_group,
+        topk_group,
+        intermediate_size,
+        local_expert_offset,
+        local_num_experts,
+        routed_scaling_factor,
+        routing_method_type,
+        use_shuffled_weight,
+        weight_layout,
+        do_finalize,
+        enable_pdl,
+        tune_max_num_tokens,
+        fp8_quantization_type,
+        activation_type,
+        norm_topk_prob,
+        routing_replay_out,
+    )
+
+    return result[0] if do_finalize else result
+
+
+def trtllm_fp8_block_scale_routed_moe(
+    topk_ids: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    do_finalize: bool = True,
+    enable_pdl: Optional[bool] = None,
+    output: Optional[torch.Tensor] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type=None,
+    activation_type: Optional[int] = None,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    from flashinfer.fused_moe.core import ActivationType, Fp8QuantizationType
+
+    if fp8_quantization_type is None:
+        fp8_quantization_type = Fp8QuantizationType.DeepSeekFp8
+    if activation_type is None:
+        activation_type = ActivationType.Swiglu.value
+    if output is None:
+        output = torch.empty(
+            hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
+        )
+
+    result = get_sgl_trtllm_moe_sm100_module().trtllm_fp8_block_scale_moe(
+        None,
+        topk_ids,
+        None,
+        routing_bias,
+        hidden_states,
+        hidden_states_scale,
+        gemm1_weights,
+        gemm1_weights_scale,
+        gemm2_weights,
+        gemm2_weights_scale,
+        output,
+        num_experts,
+        top_k,
+        n_group,
+        topk_group,
+        intermediate_size,
+        local_expert_offset,
+        local_num_experts,
+        routed_scaling_factor,
+        routing_method_type,
+        use_shuffled_weight,
+        weight_layout,
+        do_finalize,
+        enable_pdl,
+        tune_max_num_tokens,
+        fp8_quantization_type,
+        activation_type,
+        True,
+    )
+
+    return result[0] if do_finalize else result
+
+
+def trtllm_fp8_block_scale_routed_moe_lora(
+    topk_ids: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    gate_up_lora_delta: torch.Tensor,
+    activation_lora_input: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    n_group: Optional[int],
+    topk_group: Optional[int],
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    use_shuffled_weight: bool = False,
+    weight_layout: int = 0,
+    do_finalize: bool = True,
+    enable_pdl: Optional[bool] = None,
+    output: Optional[torch.Tensor] = None,
+    tune_max_num_tokens: int = 8192,
+    fp8_quantization_type=None,
+    activation_type: Optional[int] = None,
+    lora_ready_event: int = 0,
+    gemm2_done_event: int = 0,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    from flashinfer.fused_moe.core import ActivationType, Fp8QuantizationType
+    from flashinfer.utils import device_support_pdl
+
+    if fp8_quantization_type is None:
+        fp8_quantization_type = Fp8QuantizationType.DeepSeekFp8
+    if activation_type is None:
+        activation_type = ActivationType.Swiglu.value
+    if enable_pdl is None:
+        enable_pdl = device_support_pdl(hidden_states.device)
+    if output is None:
+        output = torch.empty(
+            hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
+        )
+
+    assert gate_up_lora_delta.is_contiguous()
+    assert activation_lora_input.is_contiguous()
+    empty_expert_weights = torch.empty(
+        (0,), dtype=torch.bfloat16, device=hidden_states.device
+    )
+
+    result = get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp8_block_scale_moe_lora(
+        None,
+        topk_ids,
+        empty_expert_weights,
+        routing_bias,
+        hidden_states,
+        hidden_states_scale,
+        gemm1_weights,
+        gemm1_weights_scale,
+        gemm2_weights,
+        gemm2_weights_scale,
+        output,
+        num_experts,
+        top_k,
+        n_group,
+        topk_group,
+        intermediate_size,
+        local_expert_offset,
+        local_num_experts,
+        routed_scaling_factor,
+        routing_method_type,
+        use_shuffled_weight,
+        weight_layout,
+        do_finalize,
+        enable_pdl,
+        [-1, -1],
+        fp8_quantization_type,
+        activation_type,
+        True,
+        None,
+        gate_up_lora_delta,
+        activation_lora_input,
+        lora_ready_event,
+        gemm2_done_event,
+    )
+
+    return output if do_finalize else result
+
+
+def trtllm_fp8_block_scale_moe_lora_finalize(
+    gemm2_output: torch.Tensor,
+    expert_weights: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    down_lora_delta: torch.Tensor,
+    output: torch.Tensor,
+    routed_scaling_factor: Optional[float],
+) -> torch.Tensor:
+    get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp8_block_scale_moe_lora_finalize(
+        gemm2_output,
+        expert_weights,
+        expanded_idx_to_permuted_idx,
+        down_lora_delta,
+        output,
+        routed_scaling_factor,
+    )
+    return output
+
+
+def trtllm_bf16_routed_moe_lora(
+    topk_ids: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    gemm1_weights: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gate_up_lora_delta: torch.Tensor,
+    activation_lora_input: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    do_finalize: bool = True,
+    enable_pdl: Optional[bool] = None,
+    output: Optional[torch.Tensor] = None,
+    activation_type: Optional[int] = None,
+    lora_ready_event: int = 0,
+    gemm2_done_event: int = 0,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    """BF16 MoE-LoRA: decomposed trtllm pipeline (permute -> raw gate_up GEMM ->
+    LoRA-aware activation -> down GEMM), bf16 end-to-end (no quantization).
+    Weights are the SAME prepared bf16 tensors the plain trtllm_bf16 path uses
+    (shuffled + BlockMajorK). With do_finalize=False returns
+    (gemm2_output, expert_weights, expanded_idx_to_permuted_idx) for the
+    python-side down-LoRA delta + trtllm_fp8_block_scale_moe_lora_finalize
+    (which is pure bf16 and shared across fp8/fp4/bf16)."""
+    from flashinfer.fused_moe.core import ActivationType
+    from flashinfer.utils import device_support_pdl
+
+    if activation_type is None:
+        activation_type = ActivationType.Swiglu.value
+    if enable_pdl is None:
+        enable_pdl = device_support_pdl(hidden_states.device)
+    if output is None:
+        output = torch.empty(
+            hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
+        )
+
+    assert gate_up_lora_delta.is_contiguous()
+    assert activation_lora_input.is_contiguous()
+
+    result = get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_bf16_routed_moe_lora(
+        topk_ids,
+        routing_bias,
+        hidden_states,
+        gemm1_weights,
+        gemm2_weights,
+        num_experts,
+        top_k,
+        intermediate_size,
+        local_expert_offset,
+        local_num_experts,
+        routed_scaling_factor,
+        routing_method_type,
+        do_finalize,
+        enable_pdl,
+        activation_type,
+        output,
+        True,
+        gate_up_lora_delta,
+        activation_lora_input,
+        lora_ready_event,
+        gemm2_done_event,
+    )
+
+    return output if do_finalize else result
+
+
+def trtllm_fp4_block_scale_routed_moe_lora(
+    topk_ids: torch.Tensor,
+    routing_bias: Optional[torch.Tensor],
+    hidden_states: torch.Tensor,
+    hidden_states_scale: Optional[torch.Tensor],
+    gemm1_weights: torch.Tensor,
+    gemm1_weights_scale: torch.Tensor,
+    gemm2_weights: torch.Tensor,
+    gemm2_weights_scale: torch.Tensor,
+    output1_scales_scalar: torch.Tensor,
+    output1_scales_gate_scalar: torch.Tensor,
+    output2_scales_scalar: torch.Tensor,
+    gate_up_lora_delta: torch.Tensor,
+    activation_lora_input: torch.Tensor,
+    num_experts: int,
+    top_k: int,
+    intermediate_size: int,
+    local_expert_offset: int,
+    local_num_experts: int,
+    routed_scaling_factor: Optional[float],
+    routing_method_type: int = 0,
+    do_finalize: bool = True,
+    enable_pdl: Optional[bool] = None,
+    output: Optional[torch.Tensor] = None,
+    act_type: Optional[int] = None,
+    norm_topk_prob: bool = True,
+    lora_ready_event: int = 0,
+    gemm2_done_event: int = 0,
+) -> Union[List[torch.Tensor], torch.Tensor]:
+    """NVFP4 sibling of :func:`trtllm_fp8_block_scale_routed_moe_lora`.
+
+    Decomposed (unfused-activation) MoE-LoRA on the flashinfer-trtllm backend:
+    routing -> gather -> gate_up grouped GEMM (raw 2*inter) -> activation that
+    adds ``gate_up_lora_delta`` pre-SwiGLU and writes ``activation_lora_input``
+    -> NvFP4 quant -> down grouped GEMM -> finalize. The down-LoRA is merged into
+    the output afterwards by the dispatch (virtual-experts) or via the finalize.
+
+    ``hidden_states`` is bf16 ``[num_tokens, hidden]`` (path 3: the op permutes then
+    NvFP4-quantizes it internally, ``hidden_states_scale=None``) or, legacy, packed NvFP4
+    (uint8 ``[num_tokens, hidden//2]``) with ``hidden_states_scale`` the fp8-e4m3 block
+    scale ``[num_tokens, hidden//16]``.
+    """
+    from flashinfer.fused_moe.core import ActivationType
+    from flashinfer.utils import device_support_pdl
+
+    if act_type is None:
+        act_type = ActivationType.Swiglu.value
+    if enable_pdl is None:
+        enable_pdl = device_support_pdl(hidden_states.device)
+
+    num_tokens = hidden_states.shape[0]
+    hidden_size = (
+        hidden_states.shape[1] * 2
+        if hidden_states.dtype == torch.uint8
+        else hidden_states.shape[1]
+    )
+    if output is None:
+        output = torch.empty(
+            (num_tokens, hidden_size), dtype=torch.bfloat16, device=hidden_states.device
+        )
+
+    assert gate_up_lora_delta.is_contiguous()
+    assert activation_lora_input.is_contiguous()
+    empty_expert_weights = torch.empty(
+        (0,), dtype=torch.bfloat16, device=hidden_states.device
+    )
+
+    result = get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp4_block_scale_moe_lora(
+        None,  # routing_logits (precomputed routing via packed topk_ids)
+        topk_ids,  # expert_indices (packed: (expert_id<<16)|weight_bf16.view(int16))
+        empty_expert_weights,
+        routing_bias,
+        hidden_states,
+        hidden_states_scale,
+        gemm1_weights,
+        gemm1_weights_scale,
+        None,  # gemm1_bias
+        None,  # gemm1_alpha
+        None,  # gemm1_beta
+        None,  # gemm1_clamp_limit
+        gemm2_weights,
+        gemm2_weights_scale,
+        None,  # gemm2_bias
+        output1_scales_scalar,
+        output1_scales_gate_scalar,
+        output2_scales_scalar,
+        None,  # per_token_scales (the down-GEMM act scale is produced internally)
+        num_experts,
+        top_k,
+        None,  # n_group
+        None,  # topk_group
+        intermediate_size,
+        local_expert_offset,
+        local_num_experts,
+        routed_scaling_factor,
+        routing_method_type,
+        do_finalize,
+        enable_pdl,
+        act_type,
+        output,
+        [-1, -1],  # config_index (autotuner)
+        norm_topk_prob,
+        None,  # routing_replay_out
+        gate_up_lora_delta,
+        activation_lora_input,
+        lora_ready_event,
+        bool(lora_envs.SGLANG_OPT_FUSED_PERMUTE_QUANT.get()),
+        gemm2_done_event,
+    )
+
+    return output if do_finalize else result
+
+
+def trtllm_fp4_block_scale_moe_lora_finalize(
+    gemm2_output: torch.Tensor,
+    expert_weights: torch.Tensor,
+    expanded_idx_to_permuted_idx: torch.Tensor,
+    down_lora_delta: torch.Tensor,
+    output: torch.Tensor,
+    routed_scaling_factor: Optional[float],
+) -> torch.Tensor:
+    """NvFP4 analog of :func:`trtllm_fp8_block_scale_moe_lora_finalize` — combines
+    the permuted (bf16) GEMM2 output by expert weight and merges the down-LoRA
+    delta into the per-token output (non-virtual-experts hook path)."""
+    get_sgl_trtllm_moe_sm100_raw_module().sgl_trtllm_fp4_block_scale_moe_lora_finalize(
+        gemm2_output,
+        expert_weights,
+        expanded_idx_to_permuted_idx,
+        down_lora_delta,
+        output,
+        routed_scaling_factor,
+    )
+    return output

python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_activation_quant.cuh

@@ -0,0 +1,230 @@
+// Fused SwiGLU+LoRA activation -> NVFP4 per-token quant for the FP4 MoE LoRA path.
+//
+// Modeled on tensorrt_llm::kernels::nvfp4QuantAndPerTokenScaleKernel (flashinfer
+// quantization.cuh): one block per expanded row, per-token amax via cub::BlockReduce,
+// cvt_warp_fp16_to_fp4 for the e4m3 block scale + e2m1 packing + swizzled SF layout.
+// The ONLY change vs that kernel is the pass-1 input: instead of reading activated_bf16
+// from gmem, it reads the interleaved gate/up GEMM1 output + the LoRA delta and computes
+// silu(up)*gate on the fly, rounds to bf16 (matching the standalone activation kernel's
+// bf16 output exactly), caches in smem, and also writes activation_lora_input. The down
+// GEMM input (activated_bf16) is therefore never materialized to HBM.
+//
+// Because the activated value is rounded to bf16 before quantization (same as the separate
+// activation kernel) and silu matches, the fp4 / SF / per_token_sf / activation_lora_input
+// outputs are BITWISE-identical to the unfused activation -> quant#2 chain.
+#pragma once
+
+#include <cub/block/block_reduce.cuh>
+#include <cuda/functional>
+
+#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <optional>
+
+namespace flashinfer {
+namespace sgl_fused_act_quant {
+
+namespace tk = tensorrt_llm::kernels;
+
+// Same silu as moe::dev::activation (trtllm_fused_moe_dev_kernel.cu:55): x / (1 + exp(-x)).
+inline __device__ float fused_silu(float x) {
+  return x / (1.0f + expf(-x));
+}
+
+// One block per expanded row. gateUp is the column-interleaved GEMM1 output (g0,u0,g1,u1,...)
+// indexed by permutedIdx; loraDelta is the contiguous [gate|up] delta indexed by expandedIdx.
+template <uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT, bool DISABLE_FP4_FAST_MATH>
+__global__ void fusedActivationQuantKernel(
+    int m,                                        // numTokens * topK (number of expanded rows)
+    int innerHalf,                                // inter == n (output width per row); must be a multiple of 16
+    int innerDim,                                 // gate_up_n == 2 * innerHalf
+    __nv_bfloat16 const* __restrict__ gateUp,     // interleaved gate/up, [.., innerDim] by permutedIdx
+    __nv_bfloat16 const* __restrict__ loraDelta,  // [.., innerDim] by expandedIdx, may be null
+    __nv_bfloat16* __restrict__ loraInputOut,     // [.., innerHalf] by expandedIdx, may be null
+    int32_t const* __restrict__ expandedIdxToPermutedIdx,
+    float globalScaleInv,
+    uint8_t* __restrict__ weightOutput,  // fp4 [.., innerHalf/2] by permutedIdx
+    uint8_t* __restrict__ scaleOutput,   // swizzled e4m3 SF
+    float* __restrict__ perTokenScaleOutput) {
+  constexpr int SF_VEC_SIZE = 16;
+  using InType = tk::PackedVec<__nv_bfloat16, SF_VEC_SIZE>;  // 16 bf16 == 8 __nv_bfloat162
+  using PackedFp4Type = uint64_t;                            // SF_VEC_SIZE == 16
+
+  int const expandedIdx = blockIdx.x;
+  if (expandedIdx >= m) return;
+  int const permutedIdx = expandedIdxToPermutedIdx[expandedIdx];
+  int const num_vecs_per_row = innerHalf / SF_VEC_SIZE;
+  int64_t const liBaseRow = (int64_t)expandedIdx * innerHalf;
+
+  // Padding row: the separate activation kernel writes 0 to activation_lora_input and skips
+  // the quant outputs. Mirror that, then return.
+  if (permutedIdx < 0) {
+    if (loraInputOut != nullptr) {
+      InType z;
+#pragma unroll
+      for (int i = 0; i < SF_VEC_SIZE / 2; ++i)
+        z.elts[i] = __float2bfloat162_rn(0.0f);
+      for (int vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+        *reinterpret_cast<InType*>(&loraInputOut[liBaseRow + (int64_t)vecIdx * SF_VEC_SIZE]) = z;
+      }
+    }
+    return;
+  }
+
+  int64_t const permBase = (int64_t)permutedIdx * innerDim;  // gate_up row (interleaved)
+  int64_t const expBase = (int64_t)expandedIdx * innerDim;   // delta row (contiguous gate|up)
+  (void)DISABLE_FP4_FAST_MATH;
+
+  // 1 SF block (16 outputs) per thread, held in registers across the amax barrier (no smem cache):
+  // requires num_vecs_per_row <= BLOCK_SIZE (inter=2048 -> 128 == BLOCK_SIZE). With
+  // CVT_ELTS_PER_THREAD == SF_VEC_SIZE the cvt needs no cross-thread shuffle, so masking is safe.
+  int const vecIdx = threadIdx.x;
+  bool const active = vecIdx < num_vecs_per_row;
+
+  InType vec;
+  float localAmax = 0.f;
+  if (active) {
+    int const h0 = vecIdx * SF_VEC_SIZE;
+    __nv_bfloat16 const* g = gateUp + permBase + (int64_t)2 * h0;     // 32 interleaved bf16
+    __nv_bfloat16 const* dlo = loraDelta + expBase + h0;              // silu-arg delta (lower half)
+    __nv_bfloat16 const* dhi = loraDelta + expBase + innerHalf + h0;  // multiplier delta (upper half)
+    __nv_bfloat162 amax2 = __float2bfloat162_rn(0.0f);
+    union {
+      int4 v[4];
+      __nv_bfloat16 b[32];
+    } gu;
+    union {
+      int4 v[2];
+      __nv_bfloat16 b[16];
+    } dl, dh;
+    int4 const* gp = reinterpret_cast<int4 const*>(g);
+#pragma unroll
+    for (int k = 0; k < 4; ++k)
+      gu.v[k] = gp[k];
+    if (loraDelta != nullptr) {
+      int4 const* dlp = reinterpret_cast<int4 const*>(dlo);
+      int4 const* dhp = reinterpret_cast<int4 const*>(dhi);
+#pragma unroll
+      for (int k = 0; k < 2; ++k) {
+        dl.v[k] = dlp[k];
+        dh.v[k] = dhp[k];
+      }
+    }
+#pragma unroll
+    for (int i = 0; i < SF_VEC_SIZE / 2; ++i) {  // 8 bf162 = 16 output elements
+      int const j0 = 2 * i, j1 = 2 * i + 1;
+      float even0 = (float)gu.b[2 * j0], odd0 = (float)gu.b[2 * j0 + 1];
+      float even1 = (float)gu.b[2 * j1], odd1 = (float)gu.b[2 * j1 + 1];
+      float a0 = odd0, b0 = even0, a1 = odd1, b1 = even1;
+      if (loraDelta != nullptr) {
+        a0 += (float)dl.b[j0];
+        b0 += (float)dh.b[j0];
+        a1 += (float)dl.b[j1];
+        b1 += (float)dh.b[j1];
+      }
+      float act0 = fused_silu(a0) * b0;
+      float act1 = fused_silu(a1) * b1;
+      __nv_bfloat162 e = __float22bfloat162_rn(make_float2(act0, act1));
+      vec.elts[i] = e;
+      amax2 = __hmax2(amax2, __habs2(e));
+    }
+    localAmax = (float)__hmax(amax2.x, amax2.y);
+    if (loraInputOut != nullptr) {
+      *reinterpret_cast<InType*>(&loraInputOut[liBaseRow + h0]) = vec;
+    }
+  }
+
+  // ---- per-token scale: blockReduce amax, broadcast via smem (no gmem round-trip) ----
+  using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  __shared__ float sScale;
+  float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
+  if (threadIdx.x == 0) {
+    float const pts = globalAmax * globalScaleInv;
+    perTokenScaleOutput[permutedIdx] = pts;
+    sScale = pts;
+  }
+  __syncthreads();
+  float const globalEncodeScale = tk::reciprocal_approximate_ftz(sScale);
+
+  // ---- quantize from registers (cvt computes the per-16 e4m3 block scale internally) ----
+  if (active) {
+    uint8_t fp8Scale;
+    // 5 template args on this flashinfer build: Type, SF_VEC_SIZE, CVT_ELTS_PER_THREAD,
+    // UE8M0_SF=false, TE_EXACT_NVFP4=false (the default nvfp4 quant path).
+    auto fp4Vals = tk::cvt_warp_fp16_to_fp4<__nv_bfloat16, SF_VEC_SIZE, SF_VEC_SIZE, false, false>(
+        vec, globalEncodeScale, &fp8Scale);
+    int64_t const vecOffset = (int64_t)permutedIdx * num_vecs_per_row + vecIdx;
+    reinterpret_cast<PackedFp4Type*>(weightOutput)[vecOffset] = fp4Vals;
+
+    // Match nvfp4QuantAndPerTokenScaleKernel exactly (it passes the kernel's `m` as numRows).
+    int64_t sfOffset;
+    if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+      sfOffset = vecOffset;
+    } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+      sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
+    } else {
+      sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, permutedIdx, vecIdx, m, num_vecs_per_row);
+    }
+    scaleOutput[sfOffset] = fp8Scale;
+  }
+}
+
+// Host launch: globalScaleInv = 1/448/6. BLOCK_SIZE must be >= innerHalf/16 (one SF block/thread).
+inline void launchFusedActivationQuant(
+    int m,
+    int innerHalf,
+    int innerDim,
+    __nv_bfloat16 const* gateUp,
+    __nv_bfloat16 const* loraDelta,
+    __nv_bfloat16* loraInputOut,
+    int32_t const* expandedIdxToPermutedIdx,
+    float globalScaleInv,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput,
+    tensorrt_llm::QuantizationSFLayout sfLayout,
+    bool disableFp4FastMath,
+    cudaStream_t stream) {
+  constexpr uint32_t BLOCK_SIZE = 128;  // == innerHalf/16 for inter=2048 (one SF block per thread)
+  dim3 const grid(m), block(BLOCK_SIZE);
+
+  auto launch = [&](auto layoutTag, auto fastMathTag) {
+    fusedActivationQuantKernel<BLOCK_SIZE, decltype(layoutTag)::value, decltype(fastMathTag)::value>
+        <<<grid, block, 0, stream>>>(
+            m,
+            innerHalf,
+            innerDim,
+            gateUp,
+            loraDelta,
+            loraInputOut,
+            expandedIdxToPermutedIdx,
+            globalScaleInv,
+            weightOutput,
+            scaleOutput,
+            perTokenScaleOutput);
+  };
+  auto withFastMath = [&](auto layoutTag) {
+    if (disableFp4FastMath) {
+      launch(layoutTag, std::integral_constant<bool, true>{});
+    } else {
+      launch(layoutTag, std::integral_constant<bool, false>{});
+    }
+  };
+  if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+    withFastMath(
+        std::integral_constant<
+            tensorrt_llm::QuantizationSFLayout,
+            tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
+  } else if (sfLayout == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+    withFastMath(
+        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::LINEAR>{});
+  } else {
+    withFastMath(
+        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
+  }
+}
+
+}  // namespace sgl_fused_act_quant
+}  // namespace flashinfer

python/sglang/jit_kernel/trtllm_lora_temp/data/csrc/fused_permute_quant.cuh

@@ -0,0 +1,307 @@
+// Fused permute + NvFP4-per-token-quant for the FP4 MoE-LoRA gate_up path.
+//
+// Background (decode bs64, EP8): the plain path runs `permuteKernel` (gather bf16 hidden into the
+// padded [max_padded, hidden] permuted buffer) then `nvfp4QuantAndPerTokenScaleKernel` over ALL
+// max_padded rows. At decode only num_tokens*top_k of those rows are real (the rest are padding),
+// so both kernels waste ~6x of their work on padding, and the bf16 permuted buffer is a full
+// HBM round-trip (written by permute, read back by quant).
+//
+// This kernel fuses the two: it reads the UN-permuted hidden, NvFP4-quantizes each (token,expert)
+// pair's row, and scatter-writes fp4 + swizzled block-sf + per-token-sf directly to that pair's
+// permuted position. It iterates only the num_tokens*top_k real pairs (skipping pad), and never
+// materializes the bf16 permuted buffer.
+//
+// It mirrors `nvfp4QuantAndPerTokenScaleKernel` (quantization.cuh) — same amax, same per-token-scale
+// recipe, same `cvt_warp_fp16_to_fp4`, same swizzled-sf offset (`get_sf_out_offset_8x4`) — with the
+// single `rowIdx` split into a READ row (the unpermuted source token) and a WRITE row (the permuted
+// destination). For the valid rows the result is BITWISE-identical to the plain permute->quant chain
+// (the chain's quant reads permuted_hidden[writeRow], filled by permute from hidden[readRow]; we
+// read hidden[readRow] directly), verified by the bench's fused-vs-old guard.
+//
+// PER-TOKEN-SCALE BRANCH: uses TE_EXACT (globalEncodeScale = __fdiv_rn(globalScale, globalAmax),
+// stored scale = 1/globalEncodeScale, cvt TE_EXACT_NVFP4=true). This matches the installed
+// flashinfer 0.6.11.post1, whose DISPATCH macro hard-codes the bf16 kernel to TE_EXACT_NVFP4=true
+// (quantization.cu DISPATCH_NVP4_QUANT_AND_PER_TOKEN_SCALE_KERNEL). NOTE: if a future flashinfer
+// reverts the bf16 path to the fast-math branch (reciprocal_approximate / TE_EXACT=false), or if
+// FLASHINFER_NVFP4_4OVER6 is enabled, this fused kernel would diverge — re-validate the bench's
+// fused-vs-old bitwise guard against the deployed flashinfer before trusting it there.
+//
+// Two variants (both kept, selectable, for cross-scenario perf comparison):
+//   - no-dedup: grid over the num_tokens*top_k pairs; each block re-reads+re-quantizes its source
+//     token and writes 1 destination (more blocks -> better occupancy at tiny decode sizes).
+//   - dedup:    grid over num_tokens; each block reads+quantizes its token once and scatter-writes
+//     to all of that token's (valid) permuted destinations (no redundant quant, fewer blocks).
+//
+// Helpers are pulled from quantization_utils.cuh (cvt_warp_fp16_to_fp4 / get_sf_out_offset_* /
+// PackedVec / reciprocal_approximate_ftz) rather than quantization.cuh, because the latter pulls in
+// nv_internal/.../common/cudaUtils.h, which ODR-conflicts with the flashinfer/trtllm/common twin
+// already in trtllm_fused_moe_kernel_launcher.cu's TU. loadPackedVec lives in quantization.cuh, so
+// we do a direct aligned PackedVec load instead.
+#pragma once
+
+#include <cub/cub.cuh>
+
+#include "nv_internal/tensorrt_llm/kernels/quantization_utils.cuh"
+#include <cstdint>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <optional>
+#include <type_traits>
+
+namespace sgl_fused_permute_quant {
+
+namespace tk = tensorrt_llm::kernels;
+
+// Quantize source row `readRow` of `input` (unpermuted) and write fp4 + block-sf + per-token-sf to
+// destination row `writeRow` of the permuted outputs. `numRowsSf` is the SF buffer's row count
+// (= max_padded), matching the plain quant's `m` arg to get_sf_out_offset_*.
+template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
+__device__ __forceinline__ void fused_quant_one_row(
+    uint32_t n,
+    T const* input,
+    int readRow,
+    int writeRow,
+    int numRowsSf,
+    float globalScaleInv,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput) {
+  constexpr int SF_VEC_SIZE = 16;
+  constexpr int ELTS_PER_THREAD = 16;
+  using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
+  using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
+  uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
+  uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  InType const* inBase = reinterpret_cast<InType const*>(input);
+
+  // ---- pass 1: per-row amax over the (unpermuted) source row ----
+  float localAmax = 0.f;
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
+    std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
+#pragma unroll
+    for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
+      a = __hmax2(a, __habs2(vec_in.elts[i]));
+    }
+    localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
+  }
+
+  using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
+
+  // ---- per-token scale (TE_EXACT branch — production instantiates TE_EXACT_NVFP4=true for bf16,
+  //      quantization.cu:247): globalEncodeScale = globalScale/globalAmax (exact __fdiv_rn), stored
+  //      per-token scale = 1/globalEncodeScale. __shared__ scalar replaces the gmem round-trip
+  //      (bit-identical: an fp32 store->load doesn't change the value). ----
+  __shared__ float sEncodeScale;
+  if (threadIdx.x == 0) {
+    float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
+    float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
+    sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
+  }
+  __syncthreads();
+  float const globalEncodeScale = sEncodeScale;
+  float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
+  if (threadIdx.x == 0) perTokenScaleOutput[writeRow] = perTokenScale;
+
+  // ---- pass 2: quantize + scatter-write to the permuted destination ----
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(readRow) * num_vecs_per_row + vecIdx];
+    uint8_t fp8Scale;
+    auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
+        T,
+        SF_VEC_SIZE,
+        ELTS_PER_THREAD,
+        /*UE8M0_SF=*/false,
+        /*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
+    reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
+        fp4Vals;
+
+    int64_t sfOffset;
+    if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+      sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
+    } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+      sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+    } else {
+      sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+    }
+    scaleOutput[sfOffset] = fp8Scale;
+  }
+}
+
+// no-dedup: grid.x = num_tokens*top_k (one block per (token,expert) pair).
+template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
+__global__ void fusedPermuteNvfp4QuantKernel(
+    uint32_t numPairs,
+    uint32_t n,
+    uint32_t topK,
+    int numRowsSf,
+    T const* input,
+    float globalScaleInv,
+    int32_t const* expandedIdxToPermutedIdx,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput) {
+  uint32_t const expandedIdx = blockIdx.x;
+  if (expandedIdx >= numPairs) return;
+  int const writeRow = expandedIdxToPermutedIdx[expandedIdx];
+  if (writeRow < 0) return;
+  int const readRow = static_cast<int>(expandedIdx / topK);
+  fused_quant_one_row<T, BLOCK_SIZE, SF_LAYOUT>(
+      n, input, readRow, writeRow, numRowsSf, globalScaleInv, weightOutput, scaleOutput, perTokenScaleOutput);
+}
+
+// dedup: grid.x = num_tokens (one block per source token, scatter to its top_k destinations).
+template <typename T, uint32_t BLOCK_SIZE, tensorrt_llm::QuantizationSFLayout SF_LAYOUT>
+__global__ void fusedPermuteNvfp4QuantDedupKernel(
+    uint32_t numTokens,
+    uint32_t n,
+    uint32_t topK,
+    int numRowsSf,
+    T const* input,
+    float globalScaleInv,
+    int32_t const* expandedIdxToPermutedIdx,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput) {
+  constexpr int SF_VEC_SIZE = 16;
+  constexpr int ELTS_PER_THREAD = 16;
+  using InType = tk::PackedVec<T, ELTS_PER_THREAD>;
+  using PackedFp4Type = std::conditional_t<ELTS_PER_THREAD == 16, uint64_t, uint32_t>;
+  uint32_t const token = blockIdx.x;
+  if (token >= numTokens) return;
+  uint32_t const num_vecs_per_row = (n + ELTS_PER_THREAD - 1) / ELTS_PER_THREAD;
+  uint32_t const num_sf_vecs_per_row = (n + SF_VEC_SIZE - 1) / SF_VEC_SIZE;
+  InType const* inBase = reinterpret_cast<InType const*>(input);
+
+  // pass 1: amax over the source token row (read once).
+  float localAmax = 0.f;
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
+    std::remove_reference_t<decltype(vec_in.elts[0])> a(0.f, 0.f);
+#pragma unroll
+    for (int i = 0; i < ELTS_PER_THREAD / 2; ++i) {
+      a = __hmax2(a, __habs2(vec_in.elts[i]));
+    }
+    localAmax = fmaxf(localAmax, static_cast<float>(__hmax(a.x, a.y)));
+  }
+  using BlockReduce = cub::BlockReduce<float, BLOCK_SIZE>;
+  __shared__ typename BlockReduce::TempStorage tempStorage;
+  float const globalAmax = BlockReduce(tempStorage).Reduce(localAmax, cuda::maximum<>{});
+
+  // TE_EXACT per-token scale (matches production; see fused_quant_one_row).
+  __shared__ float sEncodeScale;
+  if (threadIdx.x == 0) {
+    float const globalScale = __fdiv_rn(1.0f, globalScaleInv);
+    float const rowEncodeScale = globalAmax != 0.0f ? fminf(__fdiv_rn(globalScale, globalAmax), FLT_MAX) : FLT_MAX;
+    sEncodeScale = rowEncodeScale != 0.0f ? rowEncodeScale : 1.0f;
+  }
+  __syncthreads();
+  float const globalEncodeScale = sEncodeScale;
+  float const perTokenScale = __fdiv_rn(1.0f, globalEncodeScale);
+
+  // per-token scale -> each (valid) destination (top_k small; first top_k threads write).
+  if (threadIdx.x < topK) {
+    int const writeRow = expandedIdxToPermutedIdx[token * topK + threadIdx.x];
+    if (writeRow >= 0) perTokenScaleOutput[writeRow] = perTokenScale;
+  }
+
+  // pass 2: quantize each vec once, scatter to all valid destinations.
+  for (uint32_t vecIdx = threadIdx.x; vecIdx < num_vecs_per_row; vecIdx += BLOCK_SIZE) {
+    InType vec_in = inBase[static_cast<int64_t>(token) * num_vecs_per_row + vecIdx];
+    uint8_t fp8Scale;
+    auto fp4Vals = tk::cvt_warp_fp16_to_fp4<
+        T,
+        SF_VEC_SIZE,
+        ELTS_PER_THREAD,
+        /*UE8M0_SF=*/false,
+        /*TE_EXACT_NVFP4=*/true>(vec_in, globalEncodeScale, &fp8Scale);
+#pragma unroll 1
+    for (uint32_t k = 0; k < topK; ++k) {
+      int const writeRow = expandedIdxToPermutedIdx[token * topK + k];
+      if (writeRow < 0) continue;
+      reinterpret_cast<PackedFp4Type*>(weightOutput)[static_cast<int64_t>(writeRow) * num_vecs_per_row + vecIdx] =
+          fp4Vals;
+      int64_t sfOffset;
+      if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::LINEAR) {
+        sfOffset = static_cast<int64_t>(writeRow) * num_sf_vecs_per_row + vecIdx;
+      } else if constexpr (SF_LAYOUT == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+        sfOffset = tk::get_sf_out_offset_128x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+      } else {
+        sfOffset = tk::get_sf_out_offset_8x4(/*batchIdx=*/0, writeRow, vecIdx, numRowsSf, num_sf_vecs_per_row);
+      }
+      scaleOutput[sfOffset] = fp8Scale;
+    }
+  }
+}
+
+// Launcher. `dedup` picks the variant. `n` (= hidden) must be a multiple of 16. `numRowsSf` is the
+// SF buffer's row count (max_padded).
+template <typename T>
+void invokeFusedPermuteNvfp4Quant(
+    uint32_t numTokens,
+    uint32_t topK,
+    uint32_t n,
+    int numRowsSf,
+    T const* input,
+    float globalScaleInv,
+    int32_t const* expandedIdxToPermutedIdx,
+    uint8_t* weightOutput,
+    uint8_t* scaleOutput,
+    float* perTokenScaleOutput,
+    tensorrt_llm::QuantizationSFLayout sfLayout,
+    bool dedup,
+    cudaStream_t stream) {
+  // [opt] Occupancy tuning (ncu: kernel is occupancy-bound, not DRAM-bound — DRAM <1%, achieved
+  // occupancy was 19.8% no-dedup / 5.5% dedup at BLOCK_SIZE=128). The dedup variant launches only
+  // num_tokens CTAs (=64 at decode bs64), so it is the most CTA-starved; widening the block raises
+  // threads/CTA and hides the per-row amax-reduction + scatter latency. Decode bs64 dedup sweep:
+  // 128 -> 5.52us, 256 -> 4.09us, 512 -> 3.71us. 512 is the chosen default (the prod path uses
+  // dedup). (7168/16 = 448 vecs/row, so >448 threads idle on the tail, but the win dominates.)
+  constexpr uint32_t BLOCK_SIZE = 512;
+  dim3 const block(BLOCK_SIZE);
+
+  auto dispatch = [&](auto layoutTag) {
+    constexpr tensorrt_llm::QuantizationSFLayout LAYOUT = decltype(layoutTag)::value;
+    if (dedup) {
+      dim3 const grid(numTokens);
+      fusedPermuteNvfp4QuantDedupKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
+          numTokens,
+          n,
+          topK,
+          numRowsSf,
+          input,
+          globalScaleInv,
+          expandedIdxToPermutedIdx,
+          weightOutput,
+          scaleOutput,
+          perTokenScaleOutput);
+    } else {
+      dim3 const grid(numTokens * topK);
+      fusedPermuteNvfp4QuantKernel<T, BLOCK_SIZE, LAYOUT><<<grid, block, 0, stream>>>(
+          numTokens * topK,
+          n,
+          topK,
+          numRowsSf,
+          input,
+          globalScaleInv,
+          expandedIdxToPermutedIdx,
+          weightOutput,
+          scaleOutput,
+          perTokenScaleOutput);
+    }
+  };
+
+  if (sfLayout == tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4) {
+    dispatch(
+        std::integral_constant<
+            tensorrt_llm::QuantizationSFLayout,
+            tensorrt_llm::QuantizationSFLayout::SWIZZLED_128x4>{});
+  } else {
+    dispatch(
+        std::integral_constant<tensorrt_llm::QuantizationSFLayout, tensorrt_llm::QuantizationSFLayout::SWIZZLED_8x4>{});
+  }
+}
+
+}  // namespace sgl_fused_permute_quant

python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/DevKernel.h

@@ -0,0 +1,472 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
+#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h"
+#include <cuda.h>
+#include <iostream>
+#include <string>
+// #include <cuda_runtime_api.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_size.h>
+#include <cutlass/numeric_types.h>
+
+#include "flashinfer/exception.h"
+// #include <tensorrt_llm/common/assert.h>
+#include "flashinfer/trtllm/common/cudaUtils.h"
+#include "tensorrt_llm/common/logger.h"
+
+namespace moe::dev {
+
+#define CHECK_CUDA_ERROR(cmd)                                                                \
+  do {                                                                                       \
+    cudaError_t e = cmd;                                                                     \
+    if (e != cudaSuccess) {                                                                  \
+      std::cout << "CUDA error in " << __FILE__ << ":" << __LINE__ << " executing '" << #cmd \
+                << "': " << cudaGetErrorString(e);                                           \
+    }                                                                                        \
+    FLASHINFER_CHECK(e == cudaSuccess, "Got CUDA error. See above for details.");            \
+  } while (0)
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define LAUNCH_ESC(...) __VA_ARGS__
+
+#define LAUNCH_PDL(data, coopLaunch, types, kernel, numBlocks, numThreads, smemSize, stream)                      \
+  cudaLaunchConfig_t config{};                                                                                    \
+  config.gridDim = numBlocks;                                                                                     \
+  config.blockDim = numThreads;                                                                                   \
+  config.dynamicSmemBytes = smemSize;                                                                             \
+  config.stream = (cudaStream_t)stream;                                                                           \
+                                                                                                                  \
+  cudaLaunchAttribute attributes[2] = {};                                                                         \
+  attributes[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;                                          \
+  attributes[0].val.programmaticStreamSerializationAllowed = int(data.mUsePdl);                                   \
+  attributes[1].id = cudaLaunchAttributeCooperative;                                                              \
+  attributes[1].val.cooperative = int(coopLaunch);                                                                \
+  config.attrs = attributes;                                                                                      \
+  config.numAttrs = 2;                                                                                            \
+  if (data.mUsePdl) {                                                                                             \
+    auto params = KernelParams<types, true>::setKernelParams(data);                                               \
+    auto kernelTyped = kernel<KernelParams<types, true>>;                                                         \
+    if (smemSize > 48 * 1024)                                                                                     \
+      CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
+    CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params));                                           \
+  } else {                                                                                                        \
+    auto params = KernelParams<types, false>::setKernelParams(data);                                              \
+    auto kernelTyped = kernel<KernelParams<types, false>>;                                                        \
+    if (smemSize > 48 * 1024)                                                                                     \
+      CHECK_CUDA_ERROR(cudaFuncSetAttribute(kernelTyped, cudaFuncAttributeMaxDynamicSharedMemorySize, smemSize)); \
+    CHECK_CUDA_ERROR(cudaLaunchKernelEx(&config, kernelTyped, params));                                           \
+  }
+
+#define LAUNCH(data, kernel, numBlocks, numThreads, smemSize, stream)                                \
+  if (data.mDtypeElt == tg::Dtype::Fp16) {                                                           \
+    LAUNCH_PDL(data, false, cutlass::half_t, kernel, numBlocks, numThreads, smemSize, stream);       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3) {                                                    \
+    LAUNCH_PDL(data, false, cutlass::float_e4m3_t, kernel, numBlocks, numThreads, smemSize, stream); \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16) {                                                \
+    LAUNCH_PDL(data, false, cutlass::bfloat16_t, kernel, numBlocks, numThreads, smemSize, stream);   \
+  } else {                                                                                           \
+    FLASHINFER_WARN("Unsupported dtypeElt");                                                         \
+  }
+
+#define LAUNCH_NUM_TOKENS_PER_CTA(data, type, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream) \
+  if (numTokensPerCta == 4) {                                                                                   \
+    LAUNCH_PDL(data, false, LAUNCH_ESC(type, 4), kernel, numBlocks, numThreads, smemSize, stream);              \
+  } else if (numTokensPerCta == 2) {                                                                            \
+    LAUNCH_PDL(data, false, LAUNCH_ESC(type, 2), kernel, numBlocks, numThreads, smemSize, stream);              \
+  } else if (numTokensPerCta == 1) {                                                                            \
+    LAUNCH_PDL(data, false, LAUNCH_ESC(type, 1), kernel, numBlocks, numThreads, smemSize, stream);              \
+  } else {                                                                                                      \
+    FLASHINFER_WARN("Unsupported numTokensPerCta");                                                             \
+  }
+
+#define LAUNCH_ACTIVATION(data, kernel, numTokensPerCta, numBlocks, numThreads, smemSize, stream)       \
+  if (data.mDtypeElt == tg::Dtype::Fp16) {                                                              \
+    LAUNCH_NUM_TOKENS_PER_CTA(                                                                          \
+        data, cutlass::half_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream);       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3) {                                                       \
+    LAUNCH_NUM_TOKENS_PER_CTA(                                                                          \
+        data, cutlass::float_e4m3_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream); \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16) {                                                   \
+    LAUNCH_NUM_TOKENS_PER_CTA(                                                                          \
+        data, cutlass::bfloat16_t, numTokensPerCta, kernel, numBlocks, numThreads, smemSize, stream);   \
+  } else {                                                                                              \
+    FLASHINFER_WARN("Unsupported dtypeElt");                                                            \
+  }
+
+#define LAUNCH_EXPW(data, kernel, topK, numBlocks, numThreads, smemSize, stream)                                       \
+  if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Fp32) {                                       \
+    LAUNCH_PDL(                                                                                                        \
+        data, false, LAUNCH_ESC(cutlass::half_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream);       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Fp32) {                                \
+    LAUNCH_PDL(                                                                                                        \
+        data, false, LAUNCH_ESC(cutlass::float_e4m3_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream); \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Fp32) {                            \
+    LAUNCH_PDL(                                                                                                        \
+        data, false, LAUNCH_ESC(cutlass::bfloat16_t, float, topK), kernel, numBlocks, numThreads, smemSize, stream);   \
+  } else if (data.mDtypeElt == tg::Dtype::Fp16 && data.mDtypeExpW == tg::Dtype::Bfloat16) {                            \
+    LAUNCH_PDL(                                                                                                        \
+        data,                                                                                                          \
+        false,                                                                                                         \
+        LAUNCH_ESC(cutlass::half_t, cutlass::bfloat16_t, topK),                                                        \
+        kernel,                                                                                                        \
+        numBlocks,                                                                                                     \
+        numThreads,                                                                                                    \
+        smemSize,                                                                                                      \
+        stream);                                                                                                       \
+  } else if (data.mDtypeElt == tg::Dtype::E4m3 && data.mDtypeExpW == tg::Dtype::Bfloat16) {                            \
+    LAUNCH_PDL(                                                                                                        \
+        data,                                                                                                          \
+        false,                                                                                                         \
+        LAUNCH_ESC(cutlass::float_e4m3_t, cutlass::bfloat16_t, topK),                                                  \
+        kernel,                                                                                                        \
+        numBlocks,                                                                                                     \
+        numThreads,                                                                                                    \
+        smemSize,                                                                                                      \
+        stream);                                                                                                       \
+  } else if (data.mDtypeElt == tg::Dtype::Bfloat16 && data.mDtypeExpW == tg::Dtype::Bfloat16) {                        \
+    LAUNCH_PDL(                                                                                                        \
+        data,                                                                                                          \
+        false,                                                                                                         \
+        LAUNCH_ESC(cutlass::bfloat16_t, cutlass::bfloat16_t, topK),                                                    \
+        kernel,                                                                                                        \
+        numBlocks,                                                                                                     \
+        numThreads,                                                                                                    \
+        smemSize,                                                                                                      \
+        stream);                                                                                                       \
+  } else {                                                                                                             \
+    FLASHINFER_WARN("Unsupported pair");                                                                               \
+  }
+
+#define LAUNCH_TOPK_EXPW(data, kernel, numBlocks, numThreads, smemSize, stream) \
+  if (data.topK % 4 == 0) {                                                     \
+    LAUNCH_EXPW(data, kernel, 4, numBlocks, numThreads, smemSize, stream);      \
+  } else if (data.topK % 2 == 0) {                                              \
+    LAUNCH_EXPW(data, kernel, 2, numBlocks, numThreads, smemSize, stream);      \
+  } else {                                                                      \
+    LAUNCH_EXPW(data, kernel, 1, numBlocks, numThreads, smemSize, stream);      \
+  }
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// NOTE: Old routing-specific macros (LAUNCH_TILEN, LAUNCH_ROUTING_LLAMA4,
+// LAUNCH_ROUTING_DEEPSEEK_*, LAUNCH_ROUTING_WITH_NUM_EXPERTS) have been moved to
+// RoutingDevKernel.h which uses the new template signature with runtime isPow2/UsePdl.
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace activation {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  tg::Dtype mDtypeElt{tg::Dtype::Fp16};
+  bool mUsePdl{false};
+  bool mUseDeepSeekFp8{false};
+
+  void* inPtr;
+  void* outPtr;
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+  cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
+  cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
+
+  // When true, inPtr holds the column-interleaved gate/up GEMM1 output
+  // (g0,u0,g1,u1,...) and the kernel de-interleaves on read (x1=col 2k, x2=col 2k+1);
+  // when false, inPtr is the contiguous [gate | up] layout. Default false keeps the
+  // FP8 / non-LoRA semantics; the FP4 LoRA path sets it true to fuse the standalone
+  // de-interleave kernel into this activation read.
+  bool interleavedGateUpInput = false;
+
+  int32_t innerDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t const* totalNumPaddedTokens;
+
+  // Bench/tuning knob: when > 0, overrides the activation launch grid.x (default
+  // innerDim/128). The kernel's grid-stride hidden-dim loop makes any grid.x
+  // produce bitwise-identical output, so a smaller grid.x removes empty blocks and
+  // gives each thread a longer strip (more in-flight loads) without changing math.
+  int32_t actGridXOverride = 0;
+
+  // Activation kernel variant: 0 = scalar activationKernel, 1 = vectorized
+  // activationKernelOpt (128-bit gate/up + 64-bit delta/store, 4 pairs/thread).
+  // Variant 1 applies only to the bf16 interleaved path with (innerDim/2)%4==0;
+  // run() falls back to the scalar kernel otherwise. Output is bitwise-identical.
+  int32_t actOptMode = 0;
+};
+
+template <typename Type_, int32_t NumTokensPerCta_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  static constexpr int32_t NumTokensPerCta = NumTokensPerCta_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  Type const* inPtr;
+  Type* outPtr;
+
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+  cutlass::bfloat16_t const* gateUpLoraDeltaPtr = nullptr;
+  cutlass::bfloat16_t* activationLoraInputOutPtr = nullptr;
+
+  bool interleavedGateUpInput = false;
+
+  int32_t innerDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t const* totalNumPaddedTokens;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inPtr = (Type*)data.inPtr;
+    params.outPtr = (Type*)data.outPtr;
+    params.inDqSfsPtr = data.inDqSfsPtr;
+    params.outDqSfsPtr = data.outDqSfsPtr;
+    params.gateUpLoraDeltaPtr = data.gateUpLoraDeltaPtr;
+    params.activationLoraInputOutPtr = data.activationLoraInputOutPtr;
+    params.interleavedGateUpInput = data.interleavedGateUpInput;
+
+    params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
+
+    params.innerDim = data.innerDim;
+    params.numTokens = data.numTokens;
+    params.topK = data.topK;
+    params.totalNumPaddedTokens = data.totalNumPaddedTokens;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+}  // namespace activation
+
+namespace convertsf {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  bool mUsePdl{false};
+
+  void* inSfPtr = nullptr;
+  void* outSfPtr = nullptr;
+  int32_t hiddenDimSf;
+  int32_t numTokens;
+  tg::SfLayout sfLayoutSrc;
+  tg::SfLayout sfLayoutDst;
+};
+
+template <typename Type_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  void const* inSfPtr = nullptr;
+  void* outSfPtr = nullptr;
+  int32_t hiddenDimSf;
+  int32_t numTokens;
+  tg::SfLayout sfLayoutSrc;
+  tg::SfLayout sfLayoutDst;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inSfPtr = data.inSfPtr;
+    params.outSfPtr = data.outSfPtr;
+    params.hiddenDimSf = data.hiddenDimSf;
+    params.numTokens = data.numTokens;
+    params.sfLayoutSrc = data.sfLayoutSrc;
+    params.sfLayoutDst = data.sfLayoutDst;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace convertsf
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace permute {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  tg::Dtype mDtypeElt{tg::Dtype::Fp16};
+  bool mUsePdl{false};
+  bool mUseDeepSeekFp8{false};
+
+  void* inPtr;
+  void* outPtr;
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+  int32_t* expandedIdxToPermutedIdx;
+  int32_t hiddenDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t const* totalNumPaddedTokens;
+};
+
+template <typename Type_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  Type const* inPtr;
+  Type* outPtr;
+  float const* inDqSfsPtr;
+  float* outDqSfsPtr;
+  int32_t* expandedIdxToPermutedIdx;
+  int32_t hiddenDim;
+  int32_t numTokens;
+  int32_t topK;
+  int32_t const* totalNumPaddedTokens;
+  bool useDeepSeekFp8;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inPtr = (Type*)data.inPtr;
+    params.outPtr = (Type*)data.outPtr;
+    params.inDqSfsPtr = data.inDqSfsPtr;
+    params.outDqSfsPtr = data.outDqSfsPtr;
+    params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
+    params.hiddenDim = data.hiddenDim;
+    params.numTokens = data.numTokens;
+    params.topK = data.topK;
+    params.totalNumPaddedTokens = data.totalNumPaddedTokens;
+    params.useDeepSeekFp8 = data.mUseDeepSeekFp8;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace permute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace finalize {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace tg = batchedGemm::trtllm::gen;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Data {
+  tg::Dtype mDtypeElt{tg::Dtype::Fp16};
+  tg::Dtype mDtypeExpW{tg::Dtype::Bfloat16};
+  bool mUsePdl{false};
+  bool mUseDeepSeekFp8{false};
+
+  void* inPtr;
+  void* outPtr;
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+
+  void* expertWeightsPtr;
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t numTokens;
+  int32_t numExperts;
+  int32_t topK;
+  // Hidden dimension output of MoE block. It is not padded.
+  int32_t hiddenDim;
+  // Hidden dimension output of FC2. It might be padded.
+  int32_t hiddenDimPadded;
+  int32_t const* totalNumPaddedTokens;
+};
+
+template <typename Type_, typename TypeExpW_, int TopKUnrollFactor_, bool UsePdl_>
+struct KernelParams {
+  using Type = Type_;
+  using TypeExpW = TypeExpW_;
+  static constexpr int TopKUnrollFactor = TopKUnrollFactor_;
+  static constexpr bool UsePdl = UsePdl_;
+
+  Type const* inPtr;
+  TypeExpW const* expertWeightsPtr;
+  Type* outPtr;
+
+  float* inDqSfsPtr = nullptr;
+  float* outDqSfsPtr = nullptr;
+
+  int32_t* expandedIdxToPermutedIdx;
+
+  int32_t hiddenDim;
+  int32_t hiddenDimPadded;
+  int32_t numTokens;
+  int32_t numExperts;
+  int32_t topK;
+  int32_t const* totalNumPaddedTokens;
+
+  static KernelParams setKernelParams(Data const& data) {
+    KernelParams params;
+
+    params.inPtr = (Type*)data.inPtr;
+    params.expertWeightsPtr = (TypeExpW*)data.expertWeightsPtr;
+    params.outPtr = (Type*)data.outPtr;
+    params.inDqSfsPtr = data.inDqSfsPtr;
+    params.outDqSfsPtr = data.outDqSfsPtr;
+
+    params.expandedIdxToPermutedIdx = data.expandedIdxToPermutedIdx;
+
+    params.hiddenDim = data.hiddenDim;
+    params.hiddenDimPadded = data.hiddenDimPadded;
+    params.numTokens = data.numTokens;
+    params.numExperts = data.numExperts;
+    params.topK = data.topK;
+    params.totalNumPaddedTokens = data.totalNumPaddedTokens;
+
+    return params;
+  }
+};
+
+void run(Data const& data, void* stream);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace finalize
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace moe::dev

python/sglang/jit_kernel/trtllm_lora_temp/data/include/flashinfer/trtllm/fused_moe/runner.h

@@ -0,0 +1,586 @@
+/*
+ * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "DevKernel.h"
+#include "flashinfer/trtllm/fused_moe/RoutingKernel.h"
+#include <string>
+// #include "flashinfer/trtllm/common/cudaDriverWrapper.h"
+#include "flashinfer/trtllm/batched_gemm/KernelRunner.h"
+#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h"
+#include "flashinfer/trtllm/common/cudaUtils.h"
+
+namespace tensorrt_llm {
+namespace kernels {
+namespace trtllmgen_moe {
+
+namespace MoE {
+class Runner;
+}  // namespace MoE
+
+namespace Routing {
+
+// The type of method in top-K routing, for use in torch custom op
+// Please keep this in sync with the counterpart defined in
+// flashinfer/fused_moe/core.py
+enum class RoutingMethodType : int64_t {
+  // Default: Softmax -> TopK
+  Default = 0,
+  // Renormalize: TopK -> Softmax
+  Renormalize = 1,
+  // DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups -> Top8 experts from the
+  // Top4 groups
+  DeepSeekV3 = 2,
+  // Llama4: Top1 -> Sigmoid
+  Llama4 = 3,
+  // RenormalizeNaive: Softmax -> TopK -> Renormalize
+  RenormalizeNaive = 4,
+  // TopK only (no softmax)
+  TopK = 5,
+  // SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K weights)
+  SigmoidRenorm = 6,
+  // MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize (routeScale=1.0, epsilon=1e-20)
+  MiniMax2 = 7,
+  // Sigmoid: Sigmoid -> TopK (no renormalization)
+  Sigmoid = 8,
+  // Unspecified
+  Unspecified = 9,
+};
+
+inline int32_t maybeGetMinTokenCount(int32_t numPaddedTokens, int32_t hiddenSize, int32_t dtypeSizeBits) {
+  // Pad so total size exceeds 128KiB for performance reasons
+  int32_t minNumTokensRequired = common::divUp(128 * 1024 * 8, hiddenSize * dtypeSizeBits);
+  return std::max(numPaddedTokens, minNumTokensRequired);
+}
+
+inline std::string serializeMoeRoutingMethodType(RoutingMethodType routingMethodType) {
+  switch (routingMethodType) {
+    case RoutingMethodType::Default:
+      return "Default";
+    case RoutingMethodType::Renormalize:
+      return "Renormalize";
+    case RoutingMethodType::DeepSeekV3:
+      return "DeepSeekV3";
+    case RoutingMethodType::Llama4:
+      return "Llama4";
+    case RoutingMethodType::RenormalizeNaive:
+      return "RenormalizeNaive";
+    case RoutingMethodType::TopK:
+      return "TopK";
+    case RoutingMethodType::SigmoidRenorm:
+      return "SigmoidRenorm";
+    case RoutingMethodType::MiniMax2:
+      return "MiniMax2";
+    case RoutingMethodType::Sigmoid:
+      return "Sigmoid";
+    default:
+      return "InvalidRountingMethod";  // TODO throw error
+  };
+}
+
+inline int32_t getMaxNumCtasInBatchDim(int32_t numTokens, int32_t topK, int32_t numExperts, int32_t tileTokensDim) {
+  // For MoE, mNumTokens != 0 and the number of CTAs is known only at runtime.
+  // We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to determine
+  // the actual number of CTAs to run.
+
+  // Initialize number of tokens with the number of expanded tokens after routing.
+  int32_t numRemainingTokens = numTokens * topK;
+  int32_t maxNumCtasInBatchDim = 0;
+  // First, distribute one token each expert until token depletion to maximize CTA tile count.
+  int32_t numExpertsFilled = std::min(numExperts, numRemainingTokens);
+  maxNumCtasInBatchDim += numExpertsFilled;
+  numRemainingTokens -= numExpertsFilled;
+  // Next, greedily pour all remaining tokens to one expert to maximize CTA tile count.
+  // E.g., at this point tokens over 4 experts are [1, 1, 1, 1], and we have 4 tokens left.
+  // If each CTA handles 4 tokens/expert, the greedy strategy is to pour all remaining tokens
+  // to any one expert to get to the 5th CTA tile. Otherwise, we can only get 4 tiles in total.
+  //
+  // Another way to reason about this is to pour the remaining tokens into buckets of some fixed
+  // capacity. These buckets, if full, can then be attributed to any expert; it does not have to
+  // belong to the same expert every time.
+  if (numRemainingTokens > 0) {
+    // For every tileTokenDim tokens, we add an extra CTA tile in the token dimension.
+    // The number of CTA tiles is given by divDown(numRemainingTokens, tokenTileDim).
+    maxNumCtasInBatchDim += (numRemainingTokens / tileTokensDim);
+  }
+  return maxNumCtasInBatchDim;
+}
+
+inline int32_t
+getMaxPermutedPaddedCount(int32_t numTokens, int32_t expertsPerToken, int32_t numExperts, int32_t padding) {
+  int32_t maxCtas = getMaxNumCtasInBatchDim(numTokens, expertsPerToken, numExperts, padding);
+  return maxCtas * padding;
+}
+
+class Runner {
+ public:
+  explicit Runner();
+
+  explicit Runner(int32_t tileTokensDim);
+
+  void
+  run(void* routingLogits,
+      void* routingBias,
+      int32_t numTokens,
+      int32_t numExperts,
+      int32_t topK,
+      int32_t nGroups,
+      int32_t topkGroups,
+      int32_t localExpertOffset,
+      int32_t localNumExperts,
+      float routedScalingFactor,
+      int32_t* routingExpertIndexes,
+      int32_t* expertCountHistogram,
+      int32_t* permutedIdxSize,
+      int32_t* expandedIdxToPermutedIdx,
+      int32_t* permutedIdxToExpandedIdx,
+      int32_t* permutedIdxToTokenIdx,
+      int32_t* expertIds,
+      void* expertWeights,
+      int32_t* numTokensPerExpert,
+      int32_t* ctaIdxXyToBatchIdx,
+      int32_t* ctaIdxXyToMnLimit,
+      int32_t* numNonExitingCtas,
+      batchedGemm::trtllm::gen::Dtype dtypeElt,
+      batchedGemm::trtllm::gen::Dtype dtypeBias,
+      bool useRoutingScalesOnInput,
+      bool useDeepSeekFp8,
+      RoutingMethodType routingMethodType,
+      cudaStream_t stream,
+      batchedGemm::trtllm::gen::Dtype dtypeLogits,
+      bool normTopkProb = true,
+      int16_t* routing_replay_out = nullptr);
+
+ private:
+  friend class MoE::Runner;
+  int32_t mTileTokensDim{8};
+};
+}  // namespace Routing
+
+namespace MoE {
+// The type of activation function
+// Please keep this in sync with the counterpart defined in flashinfer/flashinfer/fused_moe/core.py
+enum class ActivationType : int64_t {
+  Gelu = 0,
+  Relu = 1,
+  Silu = 2,
+  Swiglu = 3,
+  Geglu = 4,
+  SwigluBias = 5,
+  Relu2 = 6,
+  Identity = 7,
+  InvalidType = 8,  // Must be last
+};
+
+inline std::string serializeActivationType(ActivationType activationType) {
+  switch (activationType) {
+    case ActivationType::Gelu:
+      return "Gelu";
+    case ActivationType::Relu:
+      return "Relu";
+    case ActivationType::Silu:
+      return "Silu";
+    case ActivationType::Swiglu:
+      return "Swiglu";
+    case ActivationType::Geglu:
+      return "Geglu";
+    case ActivationType::SwigluBias:
+      return "SwigluBias";
+    case ActivationType::Relu2:
+      return "Relu2";
+    case ActivationType::Identity:
+      return "Identity";
+    default:
+      return "InvalidActivationType";  // TODO throw error
+  };
+}
+
+inline bool isGatedActivation(ActivationType activationType) {
+  return activationType == ActivationType::Swiglu || activationType == ActivationType::Geglu ||
+         activationType == ActivationType::SwigluBias;
+}
+
+}  // namespace MoE
+
+namespace PermuteGemm1 {
+class Runner {
+ public:
+  explicit Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeAct,
+      batchedGemm::trtllm::gen::Dtype dtypeWeights,
+      batchedGemm::trtllm::gen::Dtype dtypeOutput,
+      bool useDeepSeekFp8,
+      int tileTokensDim,
+      MoE::ActivationType activationType,
+      bool useShuffledMatrix,
+      batchedGemm::gemm::MatrixLayout weight_layout,
+      bool usePerTokenScaling,
+      bool usePerChannelScaling,
+      bool forceUnfusedAct = false);
+
+  size_t getWorkspaceSizeInBytes(
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t configIndex) const;
+
+  [[nodiscard]] int32_t getDefaultValidConfigIndex(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
+
+  [[nodiscard]] bool isValidConfigIndex(
+      int32_t configIndex,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens) const;
+
+  [[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
+
+  void
+  run(void* hiddenState,
+      void* hiddenStateScale,
+      void* weight,
+      void* weightScale,
+      void* perTokenScales,
+      void* perChannelScales,
+      float* outputScalesScalar,
+      float* outputScalesGateScalar,
+      float* ptrBias,
+      float* ptrGatedActAlpha,
+      float* ptrGatedActBeta,
+      float* ptrClampLimit,
+      void* output,
+      void* outputScale,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t* permutedIdxToTokenIdx,
+      int32_t* ptrNumNonExitingCtas,
+      int32_t* ptrTotalNumPaddedTokens,
+      int32_t* ptrCtaIdxXyToBatchIdx,
+      int32_t* ptrCtaIdxXyToMnLimit,
+      void* bmm1Workspace,
+      bool useRoutingScalesOnInput,
+      int device,
+      cudaStream_t stream,
+      int32_t configIndex,
+      bool enable_pdl);
+
+ private:
+  friend class MoE::Runner;
+  batchedGemm::trtllm::gen::Dtype mDtypeAct;
+  batchedGemm::trtllm::gen::Dtype mDtypeWeights;
+  batchedGemm::trtllm::gen::Dtype mDtypeOutput;
+  int32_t mTileTokensDim;
+  tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
+  tensorrt_llm::kernels::trtllmgen_moe::MoE::ActivationType mActType;
+};
+}  // namespace PermuteGemm1
+
+namespace Gemm2 {
+class Runner {
+ public:
+  explicit Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeAct,
+      batchedGemm::trtllm::gen::Dtype dtypeWeights,
+      batchedGemm::trtllm::gen::Dtype outputDtype,
+      bool useDeepSeekFp8,
+      int tileTokensDim,
+      bool useShuffledMatrix,
+      batchedGemm::gemm::MatrixLayout weight_layout,
+      bool usePerTokenScaling,
+      bool usePerChannelScaling);
+
+  size_t getWorkspaceSizeInBytes(
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t configIndex) const;
+
+  [[nodiscard]] int32_t getDefaultValidConfigIndex(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numExperts, int32_t numTokens) const;
+
+  [[nodiscard]] bool isValidConfigIndex(
+      int32_t configIndex,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens) const;
+
+  [[nodiscard]] std::vector<int64_t> getPassingConfigIndices() const;
+
+  void
+  run(void* permutedHiddenState,
+      void* permutedHiddenStateScale,
+      void* weight,
+      void* weightScale,
+      void* perTokenScales,
+      void* perChannelScales,
+      float* outputScalesScalar,
+      float* ptrBias,
+      void* output,
+      void* outputScale,
+      int32_t topK,
+      int32_t hiddenSize,
+      int32_t intermediateSize,
+      int32_t numExperts,
+      int32_t numTokens,
+      int32_t* ptrNumNonExitingCtas,
+      int32_t* ptrTotalNumPaddedTokens,
+      int32_t* ptrCtaIdxXyToBatchIdx,
+      int32_t* ptrCtaIdxXyToMnLimit,
+      void* bmm2Workspace,
+      int device,
+      cudaStream_t stream,
+      int32_t configIndex,
+      bool enable_pdl);
+
+ private:
+  friend class MoE::Runner;
+  batchedGemm::trtllm::gen::Dtype mDtypeAct;
+  batchedGemm::trtllm::gen::Dtype mDtypeWeights;
+  batchedGemm::trtllm::gen::Dtype mDtypeOut;
+  int32_t mTileTokensDim;
+  tensorrt_llm::kernels::TrtllmGenBatchedGemmRunner mRunner;
+};
+}  // namespace Gemm2
+
+namespace MoE {
+namespace btg = batchedGemm::trtllm::gen;
+
+struct MoERunnerArgs {
+  void* routing_logits = nullptr;  // [num_tokens, num_experts] in float, generated after
+                                   // gemm(hidden_state, routing_weights)
+  void* routing_bias = nullptr;    // [num_experts] in bfloat16 for now = mDtypeExpW
+  void* hidden_states = nullptr;   // [num_tokens, hidden_size] in fp8 = mDtypeElt
+  // [hidden_size/128, num_tokens] in float for e4m3 DS recipe
+  // and [num_tokens, hidden_size/16] in float for e2m1
+  void* hidden_states_scale = nullptr;
+
+  // Gemm input:
+  void* gemm1_weights = nullptr;
+  void* gemm1_weights_scale = nullptr;
+  void* gemm2_weights = nullptr;
+  void* gemm2_weights_scale = nullptr;
+
+  float* gemm1_bias = nullptr;
+  float* gemm1_alpha = nullptr;
+  float* gemm1_beta = nullptr;
+  float* gemm1_clamp_limit = nullptr;
+  float* gemm2_bias = nullptr;
+
+  ActivationType activation_type = ActivationType::Swiglu;
+
+  int32_t num_tokens{0};
+  int32_t num_experts{0};
+  // Hidden dimension input of MoE block. It might be padded.
+  int32_t hidden_size{0};
+  // Hidden dimension output of MoE block. It is not padded.
+  // If not provided it is the same as hidden_size.
+  std::optional<int32_t> hidden_size_output;
+  // TODO: only compiled routing kernel supports top_k = 8
+  int32_t top_k{0};
+  int32_t n_group{0};
+  // TODO: only compiled routing kernel supports topk_group = 4
+  int32_t topk_group{0};
+  float routed_scaling_factor{0.0f};
+  int32_t intermediate_size{0};
+  int32_t local_expert_offset{0};
+  int32_t local_num_experts{0};
+  // TODO: support other types
+  btg::Dtype mDtypeElt{btg::Dtype::Void};
+  btg::Dtype mDtypeExpW{btg::Dtype::Bfloat16};
+  btg::Dtype mDtypeOut{btg::Dtype::Bfloat16};
+
+  // Apply routing scale factors to input activations
+  bool mUseRoutingScalesOnInput{false};
+  bool mUseDeepSeekFp8{false};
+  float* output1_scales_scalar = nullptr;
+  float* output1_scales_gate_scalar = nullptr;
+  float* output2_scales_scalar = nullptr;
+
+  // Optional LoRA bridge buffers used by the copied SGLang TRTLLM FP8 path.
+  // gate_up_lora_delta: [num_tokens * top_k, 2 * intermediate_size], bf16,
+  // in FlashInfer gate/up order (up first, gate second).
+  // activation_lora_input: [num_tokens * top_k, intermediate_size], bf16,
+  // populated with the post-activation intermediate for down-proj LoRA.
+  void* gate_up_lora_delta = nullptr;
+  void* activation_lora_input = nullptr;
+
+  // Optional CUDA event (cudaEvent_t) recorded on the LoRA side stream. When set, the
+  // runner waits on it right before the activation kernel (which consumes
+  // gate_up_lora_delta), so permute+GEMM1 overlap the side-stream LoRA shrink/expand
+  // instead of joining before the whole MoE op. nullptr = no wait (serial behavior).
+  void* lora_ready_event = nullptr;
+
+  // Down-LoRA/finalize overlap: optional CUDA event (cudaEvent_t) the runner records on the
+  // MoE stream right after GEMM2 (the base down GEMM), before finalize. The LoRA side stream
+  // waits on it to run the down-proj LoRA shrink/expand concurrent with the finalize kernel.
+  // nullptr = no record (serial behavior).
+  void* gemm2_done_event = nullptr;
+
+  // Output:
+  void* output = nullptr;
+  float* output_scale = nullptr;
+
+  // finalize
+  bool do_finalize{true};
+};
+
+struct MoEWorkspace {
+  // Routing intermediate outputs:
+  int32_t* routing_expert_indexes = nullptr;
+  int32_t* permuted_idx_size = nullptr;
+  int32_t* total_num_padded_tokens = nullptr;  // TODO: duplicate of permuted_idx_size
+  int32_t total_max_padded_tokens{0};
+
+  int32_t* expanded_idx_to_permuted_idx = nullptr;
+  int32_t* permuted_idx_to_expanded_idx = nullptr;
+  int32_t* permuted_idx_to_token_idx = nullptr;
+
+  // consumed by finalize kernel
+  void* expert_weights = nullptr;  // [num_tokens, top_k] in bfloat16 = mDtypeExpW
+  // consumed by permuteGemm1 kernel
+  void* token_scales = nullptr;
+  // consumed by Gemm2 kernel
+  void* token_scales_fc2 = nullptr;
+
+  int32_t* cta_idx_xy_to_batch_idx = nullptr;
+  int32_t* cta_idx_xy_to_mn_limit = nullptr;
+  int32_t* num_non_exiting_ctas = nullptr;
+
+  void* hidden_states_scale_linear = nullptr;
+
+  // Permute intermediate outputs:
+  void* permuted_hidden_states = nullptr;
+  float* permuted_hidden_states_scale = nullptr;
+
+  // Gemm1 intermediate outputs:
+  int32_t ProjUpTileN{0};
+  void* gemm1_output = nullptr;
+  float* gemm1_output_scale = nullptr;
+
+  // Activation intermediate outputs:
+  void* activation_output = nullptr;
+  float* activation_output_scale = nullptr;
+  // Unfused FP4 LoRA: bf16 [max_padded_tokens, intermediate_size] activation output written by the
+  // standalone activation kernel (gate_up LoRA added pre-SwiGLU), then NvFP4-quantized for GEMM2.
+  void* activated_lora_bf16 = nullptr;
+
+  // Gemm2 intermediate outputs:
+  void* gemm2_output = nullptr;
+  float* gemm2_output_scale = nullptr;
+
+  // Finalize intermediate outputs (placeholder not used)
+  void* finalize_output = nullptr;
+  float* finalize_output_scale = nullptr;
+
+  // FC1 workspace:
+  void* bmm1_workspace = nullptr;
+
+  // FC2 workspace:
+  void* bmm2_workspace = nullptr;
+};
+
+// Config indices to be used with Batched GEMM runners
+struct MoEConfig {
+  int64_t gemm1Config;
+  int64_t gemm2Config;
+};
+
+class Runner {
+ public:
+  // FIXME: tileTokensDim is hardcoded for now
+  Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeAct,
+      batchedGemm::trtllm::gen::Dtype dtypeWeights,
+      bool useDeepSeekFp8,
+      int tileTokensDim = 8,
+      ActivationType activationType = ActivationType::Swiglu,
+      bool useShuffledMatrix = false,
+      batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
+      bool usePerTokenScalingGemm1 = false,
+      bool usePerTokenScalingGemm2 = false,
+      bool usePerChannelScalingGemm1 = false,
+      bool usePerChannelScalingGemm2 = false,
+      bool unfuseActForLora = false);
+  Runner(
+      batchedGemm::trtllm::gen::Dtype dtypeElt,
+      bool useDeepSeekFp8,
+      int tileTokensDim = 8,
+      bool useShuffledMatrix = false,
+      batchedGemm::gemm::MatrixLayout weight_layout = batchedGemm::gemm::MatrixLayout::MajorK,
+      bool usePerTokenScalingGemm1 = false,
+      bool usePerTokenScalingGemm2 = false,
+      bool usePerChannelScalingGemm1 = false,
+      bool usePerChannelScalingGemm2 = false);
+
+  void
+  run(MoERunnerArgs const& args,
+      MoEWorkspace const& workspace,
+      int device,
+      cudaStream_t stream,
+      int64_t configIndex,
+      bool enable_pdl);
+
+  [[nodiscard]] std::tuple<int32_t, int32_t>
+  getWorkspaceSizeInBytes(MoERunnerArgs const& args, int64_t configIndex) const;
+
+  [[nodiscard]] std::vector<int64_t> getValidConfigIndices(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
+
+  [[nodiscard]] int64_t getDefaultValidConfigIndex(
+      int32_t topK, int32_t hiddenSize, int32_t intermediateSize, int32_t numLocalExperts, int32_t numTokens) const;
+
+ private:
+  void setOpsData(
+      MoERunnerArgs const& args,
+      MoEWorkspace const& workspace,
+      moe::dev::convertsf::Data& convertSfData,
+      moe::dev::activation::Data& activationData,
+      moe::dev::finalize::Data& finalizeData);
+
+ private:
+  bool mUsePerTokenScalingGemm1;
+  bool mUsePerTokenScalingGemm2;
+  bool mUsePerChannelScalingGemm1;
+  bool mUsePerChannelScalingGemm2;
+  // When true (FP4 LoRA path), GEMM1 emits the raw gate_up projection (fusedAct=false) so the
+  // standalone activation kernel can inject the gate_up LoRA delta pre-SwiGLU and capture the
+  // post-activation input for the down LoRA — mirroring the DeepSeek-FP8 unfused activation.
+  bool mUnfuseActForLora;
+  PermuteGemm1::Runner mPermuteGemm1;
+  Gemm2::Runner mGemm2;
+
+  // This will be the cartesian product of the passing configs for gemm1 and gemm2
+  // This allows us to autotune the MoE as one operation instead of tuning gemm1 and gemm2
+  // separately
+  std::vector<MoEConfig> mPassingConfigs;
+};
+}  // namespace MoE
+
+}  // namespace trtllmgen_moe
+}  // namespace kernels
+}  // namespace tensorrt_llm

python/sglang/jit_kernel/trtllm_lora_temp/jit.py

@@ -0,0 +1,101 @@
+from pathlib import Path
+
+
+def _data_dir() -> Path:
+    return Path(__file__).resolve().parent / "data"
+
+
+def gen_sgl_trtllm_gen_fused_moe_sm100_module():
+    import flashinfer
+    from flashinfer.artifacts import ArtifactPath, CheckSumHash
+    from flashinfer.jit import env as jit_env
+    from flashinfer.jit.core import current_compilation_context, gen_jit_spec
+    from flashinfer.jit.cubin_loader import (
+        ensure_symlink,
+        get_artifact,
+        get_meta_hash,
+        verify_symlinked_headers,
+    )
+    from flashinfer.jit.fused_moe import BMM_EXPORT_HEADERS
+
+    overlay_data_dir = _data_dir()
+    overlay_csrc_dir = overlay_data_dir / "csrc"
+    overlay_include_dir = overlay_data_dir / "include"
+    flashinfer_data_dir = Path(flashinfer.__file__).resolve().parent / "data"
+    flashinfer_csrc_dir = flashinfer_data_dir / "csrc"
+    flashinfer_include_dir = flashinfer_data_dir / "include"
+
+    include_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include"
+    header_name = "flashinferMetaInfo"
+    checksum_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt"
+    checksum = get_artifact(checksum_path, CheckSumHash.TRTLLM_GEN_BMM)
+    assert checksum, f"Failed to get checksums.txt from {checksum_path}"
+    meta_hash = get_meta_hash(checksum)
+
+    metainfo = get_artifact(f"{include_path}/{header_name}.h", meta_hash)
+    assert metainfo, f"{header_name}.h not found"
+
+    bmm_export_path = f"{include_path}/trtllmGen_bmm_export"
+    for header in BMM_EXPORT_HEADERS:
+        h = get_artifact(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header))
+        assert h, f"{header} not found"
+
+    symlink_path = (
+        jit_env.FLASHINFER_CUBIN_DIR
+        / "flashinfer"
+        / "trtllm"
+        / "batched_gemm"
+        / "trtllmGen_bmm_export"
+    )
+    ensure_symlink(symlink_path, jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path)
+    verify_symlinked_headers(symlink_path, BMM_EXPORT_HEADERS, checksum)
+
+    nvcc_flags = current_compilation_context.get_nvcc_flags_list(
+        supported_major_versions=[10, 12]
+    )
+
+    return gen_jit_spec(
+        "sgl_fused_moe_trtllm_sm100",
+        [
+            flashinfer_csrc_dir / "nv_internal/cpp/kernels/quantization.cu",
+            flashinfer_csrc_dir / "nv_internal/cpp/common/envUtils.cpp",
+            flashinfer_csrc_dir / "nv_internal/cpp/common/logger.cpp",
+            flashinfer_csrc_dir / "nv_internal/cpp/common/stringUtils.cpp",
+            flashinfer_csrc_dir / "nv_internal/cpp/common/tllmException.cpp",
+            flashinfer_csrc_dir / "nv_internal/cpp/common/memoryUtils.cu",
+            overlay_csrc_dir / "trtllm_fused_moe_kernel_launcher.cu",
+            overlay_csrc_dir / "trtllm_fused_moe_runner.cu",
+            flashinfer_csrc_dir
+            / "fused_moe/trtllm_backend/trtllm_fused_moe_routing_deepseek.cu",
+            flashinfer_csrc_dir
+            / "fused_moe/trtllm_backend/trtllm_fused_moe_routing_llama4.cu",
+            flashinfer_csrc_dir
+            / "fused_moe/trtllm_backend/trtllm_fused_moe_routing_custom.cu",
+            flashinfer_csrc_dir
+            / "fused_moe/trtllm_backend/trtllm_fused_moe_routing_common.cu",
+            overlay_csrc_dir
+            / "fused_moe/trtllm_backend/trtllm_fused_moe_dev_kernel.cu",
+            flashinfer_csrc_dir / "trtllm_batched_gemm_runner.cu",
+        ],
+        extra_cuda_cflags=[
+            "-DTLLM_GEN_EXPORT_INTERFACE",
+            "-DTLLM_GEN_EXPORT_FLASHINFER",
+            "-DTLLM_ENABLE_CUDA",
+            "-DENABLE_BF16",
+            "-DENABLE_FP8",
+            "-DENABLE_FP4",
+            "-DCUTLASS_ENABLE_GDC_FOR_SM100=1",
+            f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{ArtifactPath.TRTLLM_GEN_BMM}\\"',
+        ]
+        + nvcc_flags,
+        extra_include_paths=[
+            overlay_include_dir,
+            overlay_csrc_dir,
+            flashinfer_include_dir,
+            flashinfer_csrc_dir,
+            flashinfer_csrc_dir / "nv_internal",
+            flashinfer_csrc_dir / "nv_internal/include",
+            jit_env.FLASHINFER_CUBIN_DIR,
+            jit_env.FLASHINFER_CUBIN_DIR / include_path,
+        ],
+    )

python/sglang/jit_kernel/trtllm_lora_temp/kimi_k2_moe_fused_gate.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Tuple
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_kimi_k2_moe_fused_gate_module() -> Module:
+    return load_jit(
+        "kimi_k2_moe_fused_gate",
+        cuda_files=["trtllm_lora_temp/kimi_k2_moe_fused_gate.cuh"],
+        cuda_wrappers=[
+            ("kimi_k2_moe_fused_gate", "KimiK2MoEFusedGateKernel::run"),
+        ],
+    )
+
+
+def kimi_k2_moe_fused_gate(
+    input: torch.Tensor,
+    bias: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    routed_scaling_factor: float | None = 1.0,
+    apply_routed_scaling_factor_on_output: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Kimi K2 MoE fused gate (num_expert_group=1, DeepSeek noaux_tc routing).
+
+    Supports num_experts in {256, 384} and topk <= 8. input and bias are CUDA
+    tensors of float32, bfloat16, or float16 (dtypes may differ between the two);
+    they are widened to fp32 inside the kernel, so callers no longer need to
+    upcast bf16/fp16 router logits or correction bias on the host. Returns
+    (output_weights, expert_indices).
+    """
+    _supported = (torch.float32, torch.bfloat16, torch.float16)
+    assert (
+        input.dtype in _supported
+    ), f"input must be float32/bfloat16/float16, got {input.dtype}"
+    assert (
+        bias.dtype in _supported
+    ), f"bias must be float32/bfloat16/float16, got {bias.dtype}"
+    assert input.ndim == 2, "input must be 2D"
+    assert bias.ndim == 1, "bias must be 1D"
+    assert input.size(1) == bias.size(0), "input and bias must have same num_experts"
+
+    num_rows = input.size(0)
+    device = input.device
+
+    output = torch.empty(num_rows, topk, dtype=torch.float32, device=device)
+    indices = torch.empty(num_rows, topk, dtype=torch.int32, device=device)
+
+    module = _jit_kimi_k2_moe_fused_gate_module()
+    module.kimi_k2_moe_fused_gate(
+        input,
+        bias,
+        output,
+        indices,
+        topk,
+        renormalize,
+        float(routed_scaling_factor) if routed_scaling_factor is not None else 1.0,
+        apply_routed_scaling_factor_on_output,
+    )
+
+    return output, indices

python/sglang/jit_kernel/trtllm_lora_temp/moe_lora_merged_align.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "moe_lora_merged_align",
+        *args,
+        cuda_files=["trtllm_lora_temp/moe_lora_merged_align_kernel.cu"],
+        cuda_wrappers=[
+            ("moe_lora_merged_align", f"MoeLoraMergedAlignKernel<{args}>::run"),
+        ],
+    )
+
+
+def supports_merged_align(virtual_num_experts: int) -> bool:
+    """Commit-1 kernel only implements the (64, 1024] bucket-count branch.
+
+    The bucket count is virtual_num_experts + 1 (the +1 sentinel bucket). Other
+    regimes (small-batch <=64, v2 >1024) keep the old path."""
+    num_buckets = virtual_num_experts + 1
+    return 64 < num_buckets <= 1024
+
+
+def moe_lora_merged_align(
+    topk_ids: torch.Tensor,
+    token_lora_mapping: torch.Tensor,
+    num_experts: int,
+    shared_outer: bool,
+    max_loras: int,
+    block_size: int,
+    local_expert_offset: int = 0,
+    local_num_experts: Optional[int] = None,
+    do_skip: bool = True,
+    compact: bool = False,
+    fuse_scatter: Optional[bool] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    """Fused replacement for (_fused_virtual_topk_ids + _align_block_size) on the
+    merged-virtual-expert LoRA path.
+
+    Reads raw topk_ids + token_lora_mapping, computes the merged virtual id
+    inline (mirrors _fused_virtual_topk_ids), and aligns to block_size.
+
+    Returns (sorted_token_ids, expert_ids, num_tokens_post_padded,
+    token_lora_mask, virtual_num_experts).
+    """
+    device = topk_ids.device
+    flat_topk_ids = topk_ids.reshape(-1)
+    if flat_topk_ids.dtype == torch.int64:
+        flat_topk_ids = flat_topk_ids.to(torch.int32)
+    M, top_k = topk_ids.shape
+    numel = M * top_k
+
+    num_experts_for_weight = 1 if shared_outer else num_experts
+    virtual_num_experts = num_experts_for_weight * max_loras
+    ep_local = (
+        (not shared_outer)
+        and (local_num_experts is not None)
+        and (local_num_experts < num_experts_for_weight)
+    )
+
+    # compact: histogram over only the rank's local experts (dense LOCAL ids)
+    # instead of the full global virtual space. Valid only for the single-adapter
+    # EP per-expert path (safe_lora shift is 0; owned ids are a contiguous window
+    # remappable by a single -offset). expert_ids is restored to global in-kernel.
+    compact_eff = compact and ep_local and max_loras == 1 and not shared_outer
+    bucket_experts = local_num_experts if compact_eff else virtual_num_experts
+
+    # fuse_scatter: do the whole align+scatter in one threadblock (one launch).
+    # Only for small numel (the scatter is single-block); large numel (prefill)
+    # keeps the 2-kernel multi-block path. Default auto by numel.
+    fuse_eff = (numel <= 2048) if fuse_scatter is None else fuse_scatter
+    if fuse_eff:
+        # The fused kernel's dynamic shared memory must fit the 48KB default limit
+        # (no cudaFuncSetAttribute opt-in). Layout matches the kernel exactly:
+        # shared_counts + prefix + scan_buf + warp_sums + cursor + svids.
+        nb = bucket_experts + 1  # bucket count (the kernel's num_experts)
+        scan_size = 1 << (nb - 1).bit_length() if nb > 1 else 1
+        fused_shmem = (nb + (nb + 1) + scan_size + 32 + nb + numel) * 4
+        if fused_shmem > 47 * 1024:
+            fuse_eff = False  # too big -> fall back to the 2-kernel path
+
+    # Allocation mirrors moe_align_block_size.py (the kernel uses a +1 sentinel
+    # bucket, so the padded buffers are sized with bucket_experts + 1).
+    if numel < bucket_experts + 1:
+        max_num_tokens_padded = numel * block_size
+    else:
+        max_num_tokens_padded = numel + (bucket_experts + 1) * (block_size - 1)
+    max_num_m_blocks = (max_num_tokens_padded + block_size - 1) // block_size
+
+    # The align kernel's block-1 fill writes sorted_token_ids with vectorized int4
+    # stores; the last store can spill up to 3 int32 past the logical end. Pad the
+    # standalone allocation to a multiple of VEC_SIZE (4) so the spill stays in
+    # bounds (matches _align_block_size_jit's _A4). block_size=16 is already a
+    # multiple of 4 in production; this guards non-multiple-of-4 block sizes.
+    sorted_alloc = (max_num_tokens_padded + 3) & ~3
+    sorted_token_ids = torch.empty((sorted_alloc,), dtype=torch.int32, device=device)
+    expert_ids = torch.empty((max_num_m_blocks,), dtype=torch.int32, device=device)
+    num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device=device)
+    # No memset: the align kernel writes cumsum[0..num_buckets] before
+    # count_and_sort reads it (same as moe_align_block_size.py's empty cumsum).
+    cumsum_buffer = torch.empty((bucket_experts + 2,), dtype=torch.int32, device=device)
+    token_lora_mask = torch.empty((M,), dtype=torch.bool, device=device)
+
+    module = _jit_module(flat_topk_ids.dtype)
+    module.moe_lora_merged_align(
+        flat_topk_ids,
+        token_lora_mapping,
+        token_lora_mask,
+        bucket_experts + 1,  # bucket count (the +1 sentinel)
+        block_size,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        cumsum_buffer,
+        True,  # pad_sorted_token_ids
+        top_k,
+        num_experts_for_weight,
+        local_expert_offset,
+        local_num_experts if local_num_experts is not None else 0,
+        ep_local,
+        shared_outer,
+        do_skip,
+        compact_eff,
+        fuse_eff,
+    )
+
+    return (
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        token_lora_mask,
+        virtual_num_experts,
+    )

python/sglang/jit_kernel/trtllm_lora_temp/topk_pack.py

@@ -0,0 +1,51 @@
+"""Fused pack for the trtllm routed-MoE topk format.
+
+The trtllm routed MoE consumes top-k routing as a single int32 per (token, slot):
+``PackedScoreIdx`` = ``(expert_id << 16) | bf16_weight_bits`` (little-endian: low 16
+bits = bf16 weight, high 16 bits = int16 expert id).
+
+The torch reference builds this with a cluster of ~4 elementwise ops (cast, cast, view,
+bitshift, or). On the decode LoRA path these run on the main CUDA stream between
+``per_token_group_quant_fp8`` and the trtllm MoE op. This kernel collapses them into a
+single Triton launch. Bit-identical to the torch reference: weights are softmax/renormalized
+(>= 0 -> bf16 sign bit is 0, so the low 16 bits never collide with the id field and masking
+== torch's sign-extend-then-or); expert ids are small (< num_experts).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _pack_topk_kernel(ids_ptr, w_ptr, out_ptr, numel, BLOCK: tl.constexpr):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < numel
+    ids = tl.load(ids_ptr + offs, mask=mask)  # int32
+    w = tl.load(w_ptr + offs, mask=mask)  # float32
+    wb = w.to(tl.bfloat16)
+    wbits = wb.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF
+    packed = (ids << 16) | wbits
+    tl.store(out_ptr + offs, packed, mask=mask)
+
+
+def fused_pack_topk(topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor:
+    """Single-launch replacement for the elementwise routed-MoE topk pack.
+
+    Returns int32 ``[*, top_k]`` packed tensor, bit-identical to the torch reference.
+    """
+    if topk_ids.dtype != torch.int32:
+        topk_ids = topk_ids.to(torch.int32)
+    topk_ids = topk_ids.contiguous()
+    if topk_weights.dtype != torch.float32:
+        topk_weights = topk_weights.to(torch.float32)
+    topk_weights = topk_weights.contiguous()
+    out = torch.empty_like(topk_ids, dtype=torch.int32)
+    numel = out.numel()
+    if numel == 0:
+        return out
+    BLOCK = 1024
+    grid = (triton.cdiv(numel, BLOCK),)
+    _pack_topk_kernel[grid](topk_ids, topk_weights, out, numel, BLOCK=BLOCK)
+    return out

python/sglang/jit_kernel/trtllm_lora_temp/topk_softmax_pack.py

@@ -0,0 +1,90 @@
+"""Fused top-k gating softmax with routed-pack output (JIT).
+
+JIT port of sgl-kernel's AOT ``topk_softmax`` power-of-2 fast path
+(``topkGatingSoftmax``) extended with a third output: the FlashInfer routed-MoE
+packed format ``(topk_id << 16) | bf16_bits(topk_weight)`` computed in the
+kernel epilogue after renormalization — bit-identical to running the standalone
+``fused_pack_topk`` triton kernel on the post-processed topk_ids/topk_weights
+(including the ``_mask_topk_ids_padded_region`` id=-1 sentinel for rows at or
+beyond ``num_token_non_padded``). This removes the per-MoE-layer
+``_pack_topk_kernel`` launch from the decode critical path.
+
+Scope (callers must fall back to the AOT ``topk_softmax`` + separate pack
+otherwise): power-of-2 ``num_experts`` in [1, 512]; no softcapping or
+correction bias (the Qwen3-MoE softmax routing uses neither).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_topk_softmax_pack_module() -> Module:
+    return load_jit(
+        "topk_softmax_pack",
+        cuda_files=["trtllm_lora_temp/topk_softmax_pack.cuh"],
+        cuda_wrappers=[("topk_softmax_pack", "topk_softmax_pack")],
+    )
+
+
+@register_custom_op(mutates_args=["topk_weights", "topk_indices", "packed"])
+def _jit_topk_softmax_pack_op(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    packed: torch.Tensor,
+    gating_output: torch.Tensor,
+    num_token_non_padded: Optional[torch.Tensor],
+    renormalize: bool,
+) -> None:
+    module = _jit_topk_softmax_pack_module()
+    module.topk_softmax_pack(
+        topk_weights,
+        topk_indices,
+        packed,
+        gating_output,
+        num_token_non_padded,
+        renormalize,
+    )
+
+
+def topk_softmax_pack(
+    topk_weights: torch.Tensor,
+    topk_indices: torch.Tensor,
+    packed: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool,
+    num_token_non_padded: Optional[torch.Tensor] = None,
+) -> None:
+    """Drop-in for the AOT ``topk_softmax`` that ALSO writes ``packed``.
+
+    ``packed`` is int32 ``[num_tokens, topk]``, ``(id << 16) | bf16_bits(w)``
+    with the final (renormalized) weights; rows >= ``num_token_non_padded``
+    pack id = -1 (the padded-region sentinel). ``topk_weights``/``topk_indices``
+    are written exactly like the AOT kernel (indices NOT masked here — the
+    regular python post-process handles them).
+    """
+    assert gating_output.dim() == 2
+    num_experts = gating_output.shape[-1]
+    assert num_experts & (num_experts - 1) == 0 and num_experts <= 512, (
+        "topk_softmax_pack supports power-of-2 num_experts in [1, 512] only; "
+        "fall back to topk_softmax + fused_pack_topk"
+    )
+    if gating_output.shape[0] == 0:
+        return
+    _jit_topk_softmax_pack_op(
+        topk_weights,
+        topk_indices,
+        packed,
+        gating_output.contiguous(),
+        num_token_non_padded,
+        renormalize,
+    )