fix: address code review comments for kunlunxin vendor backend

Author: luoyc123

vllm_fl/dispatch/backends/vendor/kunlunxin/impl/causal_conv1d.py

@@ -118,6 +118,22 @@ def causal_conv1d_update_kunlunxin(
     Returns:
         output tensor (same shape as input x)
     """
+    # Reject speculative decoding: not supported by native kernel
+    if num_accepted_tokens is not None:
+        raise NotImplementedError(
+            "Kunlunxin causal_conv1d_update does not support speculative decoding. "
+            "The native xtorch_ops.causal_conv1d_update kernel does not accept "
+            "num_accepted_tokens parameter."
+        )
+
+    # Reject varlen mode with query_start_loc: not supported by native kernel
+    if query_start_loc is not None and max_query_len > 0:
+        raise NotImplementedError(
+            "Kunlunxin causal_conv1d_update does not support varlen with query_start_loc. "
+            "The native xtorch_ops.causal_conv1d_update kernel does not accept "
+            "query_start_loc and max_query_len parameters."
+        )
+
     if activation not in [None, "silu", "swish"]:
         raise NotImplementedError(
             f"activation must be None, silu, or swish, actual: {activation}"

vllm_fl/dispatch/backends/vendor/kunlunxin/impl/fla/chunk.py

@@ -16,8 +16,9 @@
 Contains both the low-level kernel wrapper (chunk_gated_delta_rule_fwd) and
 the top-level entry (chunk_gated_delta_rule).
 
-The top-level ChunkGatedDeltaRuleFunction.forward **skips l2norm** because the
-Kunlunxin kernel handles it internally (use_qk_l2norm_in_kernel=True).
+The Kunlunxin native kernel supports use_qk_l2norm_in_kernel natively,
+so normalization is delegated to the kernel when enabled, rather than
+being done in Python like other backends.
 """
 
 from __future__ import annotations
@@ -41,6 +42,7 @@ def chunk_gated_delta_rule_fwd(
     initial_state: torch.Tensor,
     output_final_state: bool,
     cu_seqlens: Optional[torch.LongTensor] = None,
+    use_qk_l2norm_in_kernel: bool = True,
 ) -> tuple:
     """
     Kunlunxin chunked gated delta rule forward pass.
@@ -74,6 +76,9 @@ def chunk_gated_delta_rule_fwd(
 
     cu_seqlens_cpu = cu_seqlens.cpu() if cu_seqlens is not None else None
 
+    # Use -1 for automatic scale (1/sqrt(K)) only when caller does not provide one
+    kernel_scale = scale if scale is not None else -1
+
     final_state = torch.empty_like(state_input)
     o = torch.empty_like(v)
     xtorch_ops.chunk_gated_delta_rule(
@@ -82,13 +87,13 @@ def chunk_gated_delta_rule_fwd(
         v,
         g_input,
         beta_input,
-        -1,  # uses -1 for automatic scale
+        kernel_scale,
         state_input,
         o,
         final_state,
         cu_seqlens_cpu,
         head_first=False,
-        use_qk_l2norm_in_kernel=True,  # always True in Kunlunxin
+        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
     )
 
     # Transpose final_state back to [N, H, K, V]
@@ -115,8 +120,8 @@ def forward(
         cu_seqlens: Optional[torch.LongTensor] = None,
         use_qk_l2norm_in_kernel: bool = False,
     ):
-        # klx diff: skip l2norm — kernel handles it internally
-        # (upstream would do l2norm_fwd(q), l2norm_fwd(k) here)
+        # When use_qk_l2norm_in_kernel=True, the native kernel handles L2-norm
+        # internally. When False, the caller has already normalized q/k externally.
 
         # [N, H, K, V] -> [N, H, V, K] for kunlunxin kernel
         if initial_state is not None:
@@ -132,6 +137,7 @@ def forward(
             initial_state=initial_state,
             output_final_state=output_final_state,
             cu_seqlens=cu_seqlens,
+            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
         )
         ctx.scale = scale
         ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel

vllm_fl/dispatch/backends/vendor/kunlunxin/impl/fused_moe/experts_selector.py

@@ -26,11 +26,11 @@ def vllm_topk_softmax(
 ) -> tuple[torch.Tensor, ...]:
     if renormalize:
         xtorch_ops.moe_softmax_topk_norm(
-            gating_output, topk_weights, topk_indices, None,
+            gating_output, topk_weights, topk_indices, token_expert_indices,
         )
     else:
         xtorch_ops.moe_softmax_topk(
-            gating_output, topk_weights, topk_indices, None,
+            gating_output, topk_weights, topk_indices, token_expert_indices,
         )
     return topk_weights, topk_indices
 
@@ -78,44 +78,60 @@ def fused_topk_bias(
 
 
 def grouped_topk(
-    hidden_states: torch.Tensor,
-    gating_output: torch.Tensor,
+    scores: torch.Tensor,
+    n_group: int,
+    topk_group: int,
     topk: int,
     renormalize: bool,
-    num_expert_group: int = 0,
-    topk_group: int = 0,
-    scoring_func: str = "softmax",
-    routed_scaling_factor: float = 1.0,
-    e_score_correction_bias: Optional[torch.Tensor] = None,
+    routed_scaling_factor: float,
+    bias: torch.Tensor,
+    scoring_func: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
-    """Grouped top-k selection."""
-    seq_num = gating_output.shape[0]
-
-    if scoring_func == "softmax":
-        scores = gating_output.softmax(dim=-1)
-    elif scoring_func == "sigmoid":
-        scores = gating_output.sigmoid()
-    else:
-        raise ValueError(f"Unsupported scoring_func: {scoring_func}")
-
-    if e_score_correction_bias is not None:
-        assert e_score_correction_bias.dtype == torch.float32
-        scores_for_choice = scores + e_score_correction_bias.unsqueeze(0)
+    """Grouped top-k selection with unified dispatch ABI.
+
+    Args:
+        scores: Already computed scores tensor (after softmax/sigmoid if needed)
+        n_group: Number of expert groups
+        topk_group: Number of groups to select
+        topk: Total number of experts to select
+        renormalize: Whether to renormalize weights
+        routed_scaling_factor: Scaling factor for routing weights
+        bias: Score correction bias tensor
+        scoring_func: 0=none (scores already processed), 1=sigmoid
+
+    Returns:
+        topk_weights: Selected expert weights
+        topk_ids: Selected expert indices
+    """
+    seq_num = scores.shape[0]
+
+    # Apply sigmoid if scoring_func=1 (dispatcher may pass raw logits in this case)
+    if scoring_func == 1:
+        scores = scores.sigmoid()
+
+    # Apply bias for expert selection
+    if bias is not None and bias.numel() > 0:
+        assert bias.dtype == torch.float32
+        scores_for_choice = scores + bias.unsqueeze(0)
     else:
         scores_for_choice = scores
 
-    topk_weights = torch.empty((seq_num, topk), dtype=torch.float, device=gating_output.device)
-    topk_ids = torch.empty((seq_num, topk), dtype=torch.int32, device=gating_output.device)
+    topk_weights = torch.empty((seq_num, topk), dtype=torch.float, device=scores.device)
+    topk_ids = torch.empty((seq_num, topk), dtype=torch.int32, device=scores.device)
 
     xtorch_ops.moe_group_topk(
-        scores_for_choice, num_expert_group, topk_group,
+        scores_for_choice, n_group, topk_group,
         topk_weights, topk_ids, None,
     )
 
-    if e_score_correction_bias is not None:
+    # If bias was used for selection, gather original scores for weights
+    if bias is not None and bias.numel() > 0:
         topk_weights = scores.gather(1, topk_ids)
 
     if renormalize:
         topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
 
+    if routed_scaling_factor != 1.0:
+        topk_weights = topk_weights * routed_scaling_factor
+
     return topk_weights, topk_ids

vllm_fl/dispatch/backends/vendor/kunlunxin/impl/fused_moe/fused_moe.py

@@ -30,15 +30,23 @@ def _klx_fused_experts(
     w2: torch.Tensor,
     topk_weights: torch.Tensor,
     topk_ids: torch.Tensor,
+    activation: str = "silu",
     use_int8_w8a8: bool = False,
     use_int8_w4a8: bool = False,
     w1_scale: Optional[torch.Tensor] = None,
     w2_scale: Optional[torch.Tensor] = None,
+    w1_bias: Optional[torch.Tensor] = None,
+    w2_bias: Optional[torch.Tensor] = None,
 ) -> None:
     """
     Fused MoE expert computation using xtorch_ops (sorted path).
 
-    Pipeline: gen_block_statistic -> moe_pre_sorted -> moe_fc(w1) -> swiglu -> moe_fc(w2) -> post (weight+sum)
+    Pipeline: gen_block_statistic -> moe_pre_sorted -> moe_fc(w1) -> activation -> moe_fc(w2) -> post (weight+sum)
+
+    Args:
+        activation: Activation function type. Supported: "silu", "gelu", "relu"
+        w1_bias: Optional bias for gate+up projection [E, 2*ffn_hd]
+        w2_bias: Optional bias for down projection [E, hidden_dim]
     """
     if use_int8_w8a8 or use_int8_w4a8:
         raise NotImplementedError("_klx_fused_experts is not supported for int8 w8a8 and w4a8.")
@@ -73,18 +81,43 @@ def _klx_fused_experts(
     inner_fc_out = torch.empty(seq_num, moe_topk, double_ffn_hd, dtype=dtype, device=device)
     xtorch_ops.moe_fc(
         moe_expand, w1, sorted_tokens_num_lod, moe_index, moe_topk, inner_fc_out,
+        bias=w1_bias,
     )
     inner_fc_out = inner_fc_out.view(moe_input_num, double_ffn_hd)
 
-    # Step 4: SwiGLU activation (in-place on first half)
+    # Step 4: Activation
     ffn_hd = double_ffn_hd // 2
-    swiglu_out = torch.empty(moe_input_num, ffn_hd, dtype=dtype, device=device)
-    xtorch_ops.swiglu(inner_fc_out, swiglu_out)
+    if activation == "silu":
+        # SwiGLU: silu(gate) * up
+        swiglu_out = torch.empty(moe_input_num, ffn_hd, dtype=dtype, device=device)
+        xtorch_ops.swiglu(inner_fc_out, swiglu_out)
+    elif activation == "gelu":
+        # GeGLU: gelu(gate) * up
+        gate = inner_fc_out[:, :ffn_hd]
+        up = inner_fc_out[:, ffn_hd:]
+        swiglu_out = xtorch_ops.gelu(gate) * up
+    elif activation == "relu":
+        # ReLU: relu(gate) * up
+        gate = inner_fc_out[:, :ffn_hd]
+        up = inner_fc_out[:, ffn_hd:]
+        swiglu_out = xtorch_ops.relu(gate) * up
+    elif activation in ["gelu_no_mul", "silu_no_mul"]:
+        # No mul variant: only apply activation, no gating
+        if activation == "gelu_no_mul":
+            swiglu_out = xtorch_ops.gelu(inner_fc_out)
+        else:  # silu_no_mul
+            swiglu_out = xtorch_ops.silu(inner_fc_out)
+    else:
+        raise ValueError(
+            f"Unsupported activation '{activation}'. "
+            f"Supported: ['silu', 'gelu', 'relu', 'gelu_no_mul', 'silu_no_mul']"
+        )
 
     # Step 5: Outer FC (down projection)
     outer_fc_out = torch.empty(seq_num, moe_topk, hidden_dim, dtype=dtype, device=device)
     xtorch_ops.moe_fc(
         swiglu_out, w2, sorted_tokens_num_lod, moe_index, moe_topk, outer_fc_out,
+        bias=w2_bias,
     )
     outer_fc_out = outer_fc_out.view(moe_input_num, hidden_dim)
 
@@ -127,13 +160,37 @@ def fused_experts_impl(
 ) -> torch.Tensor:
     """
     Kunlunxin fused experts implementation.
-    
+
     This function matches the signature of vllm_fl.ops.fused_moe.fused_moe.fused_experts_impl
     and is patched in by the Kunlunxin patch system.
+
+    Args:
+        activation: Activation function. Supported: "silu", "gelu", "relu", "gelu_no_mul", "silu_no_mul"
+        apply_router_weight_on_input: If True, apply router weights on input (NOT SUPPORTED)
+        w1_bias: Optional bias for gate+up projection
+        w2_bias: Optional bias for down projection
     """
+    # Stage 1: Explicit rejections for unsupported features
+
+    # 1.1: Reject unsupported quantization schemes
     if use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16:
         raise NotImplementedError(
-            "Kunlunxin fused_experts does not support fp8/int8_w8a16/int4 quantization yet."
+            "Kunlunxin fused_experts does not support fp8_w8a8/int8_w8a16/int4_w4a16 quantization yet."
+        )
+
+    # 1.2: Reject apply_router_weight_on_input=True
+    if apply_router_weight_on_input:
+        raise NotImplementedError(
+            "Kunlunxin fused_experts does not support apply_router_weight_on_input=True. "
+            "Router weights are always applied in the moe_post stage (after down projection)."
+        )
+
+    # 1.3: Validate activation function
+    SUPPORTED_ACTIVATIONS = ["silu", "gelu", "relu", "gelu_no_mul", "silu_no_mul"]
+    if activation not in SUPPORTED_ACTIVATIONS:
+        raise NotImplementedError(
+            f"Kunlunxin fused_experts does not support activation '{activation}'. "
+            f"Supported activations: {SUPPORTED_ACTIVATIONS}"
         )
 
     num_tokens = hidden_states.size(0)
@@ -157,9 +214,12 @@ def fused_experts_impl(
         w2=w2,
         topk_weights=topk_weights,
         topk_ids=topk_ids,
+        activation=activation,
         use_int8_w8a8=use_int8_w8a8,
         w1_scale=w1_scale,
         w2_scale=w2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
     )
 
     return output

vllm_fl/dispatch/backends/vendor/kunlunxin/kunlunxin.py

@@ -135,7 +135,9 @@ def attention_backend(self, use_mla: bool = False, use_sparse: bool = False) ->
         Returns:
             Fully qualified class path string
         """
+        # TODO: Implement MLA (Multi-head Latent Attention) support for Kunlunxin
         if use_mla:
+            # TODO: Implement MLA with sparse attention support for Kunlunxin
             if use_sparse:
                 raise NotImplementedError("MLA with sparse attention is not implemented for Kunlunxin yet.")
             raise NotImplementedError("MLA attention is not implemented for Kunlunxin yet.")
@@ -193,22 +195,23 @@ def topk_softmax(
 
     def grouped_topk(
         self,
-        hidden_states: 'torch.Tensor',
-        gating_output: 'torch.Tensor',
+        scores: 'torch.Tensor',
+        n_group: int,
+        topk_group: int,
         topk: int,
         renormalize: bool,
-        num_expert_group: int = 0,
-        topk_group: int = 0,
-        scoring_func: str = 'softmax',
-        routed_scaling_factor: float = 1.0,
-        e_score_correction_bias: 'torch.Tensor | None' = None,
+        routed_scaling_factor: float,
+        bias: 'torch.Tensor',
+        scoring_func: int = 0,
     ) -> tuple:
         from .impl.fused_moe.experts_selector import grouped_topk
         return grouped_topk(
-            hidden_states, gating_output, topk, renormalize,
-            num_expert_group=num_expert_group,
-            topk_group=topk_group,
-            scoring_func=scoring_func,
-            routed_scaling_factor=routed_scaling_factor,
-            e_score_correction_bias=e_score_correction_bias,
+            scores,
+            n_group,
+            topk_group,
+            topk,
+            renormalize,
+            routed_scaling_factor,
+            bias,
+            scoring_func,
         )

vllm_fl/dispatch/backends/vendor/kunlunxin/patch.py

@@ -43,6 +43,7 @@ def apply_kunlunxin_patches():
     patch_fused_gdn_gating()
     patch_ssm_cache_update()
     patch_sampler_rng()
+    patch_decode_attention()
     logger.info("Applied all Kunlunxin patches")
 
 

vllm_fl/dispatch/config/kunlunxin.yaml

@@ -33,19 +33,19 @@ op_backends:
     - vendor
     - reference
 
-  # rms_norm: prioritize vendor (Kunlunxin) implementation
+  # rms_norm: prioritize flagos (FlagGems) implementation, fallback to vendor
   rms_norm:
     - flagos
     - vendor
     - reference
 
-  # silu_and_mul: prioritize vendor (Kunlunxin) implementation
+  # silu_and_mul: prioritize flagos (FlagGems) implementation, fallback to vendor
   silu_and_mul:
     - flagos
     - vendor
     - reference
 
-  # rotary_embedding: prioritize vendor (Kunlunxin) implementation
+  # rotary_embedding: prioritize flagos (FlagGems) implementation, fallback to vendor
   rotary_embedding:
     - flagos
     - vendor