fix: group_gemm format "TN" + Layer 3 ops_api dispatch from xllm upstream

AST chain alignment with upstream_ref/xllm/xllm/core/kernels/ilu/:

Layer 5: ixformer::infer (binary .so on device)
Layer 4: xllm_kernels/ilu/*.cpp -> calls ixformer::infer (0-diff with upstream)
Layer 3: xllm_kernels/ops_api.h+cpp + param.h (NEW from upstream 2719 lines)
         kernels/kernels.h aggregation header (NEW)
Layer 2: xllm_layers/ilu/*.cpp (0-diff with upstream)
Layer 1: ix_full_bridge_v2.cpp pybind11 bridge (FIXED)

Critical fixes in ix_full_bridge_v2.cpp:
- group_gemm format "default" -> "TN" (match upstream ilu/group_gemm.cpp)
- fused_moe_forward: pass 3D weights directly instead of .view({-1,...})
- group_gemm output_n: use tokens_per_experts.sum() per upstream convention

ex_engine/csrc/ix_full_bridge_v2.cpp

@@ -326,15 +326,21 @@ torch::Tensor ix_moe_expand_input(torch::Tensor input,
 torch::Tensor ix_group_gemm(torch::Tensor inputs, torch::Tensor weights,
                             torch::Tensor tokens_per_experts,
                             int64_t output_n) {
+    // Match upstream xllm/core/kernels/ilu/group_gemm.cpp exactly:
+    // moe_w16a16_group_gemm(output, input, weight, tokens_per_experts,
+    //                       dst_to_src=nullopt, bias=nullopt,
+    //                       format="TN", persistent=0,
+    //                       output_n=tokens_per_experts.sum())
     int64_t total_tokens = inputs.size(0);
     auto output = inputs.new_empty({total_tokens, output_n});
+    int64_t gemm_output_n = tokens_per_experts.sum().item<int64_t>();
     ixformer::infer::moe_w16a16_group_gemm(
         output, inputs, weights, tokens_per_experts,
         /*dst_to_src=*/c10::nullopt,
         /*bias=*/c10::nullopt,
-        /*format=*/"default",
+        /*format=*/"TN",
         /*persistent=*/0,
-        output_n);
+        gemm_output_n);
     return output;
 }
 
@@ -382,16 +388,20 @@ torch::Tensor ix_fused_moe_forward(
     auto expanded = ix_moe_expand_input(hidden_states, src_dst, dst_src, topk);
 
     // Step 4: group_gemm (w13: gate_up projection)
+    // w13 shape: [num_experts, 2*intermediate, hidden] — pass as-is (3D)
+    // output_n = tokens_per_experts.sum() per upstream convention
     int64_t intermediate_2x = w13.size(1);
-    auto gate_up = ix_group_gemm(expanded, w13.view({-1, w13.size(2)}),
+    int64_t output_n_w13 = expert_sizes_gpu.sum().item<int64_t>();
+    auto gate_up = ix_group_gemm(expanded, w13,
                                  expert_sizes_gpu, intermediate_2x);
 
     // Step 5: silu_and_mul
     auto activated = ix_silu_and_mul(gate_up);
 
     // Step 6: group_gemm (w2: down projection)
+    // w2 shape: [num_experts, hidden, intermediate] — pass as-is (3D)
     int64_t hidden_size = w2.size(1);
-    auto down = ix_group_gemm(activated, w2.view({-1, w2.size(2)}),
+    auto down = ix_group_gemm(activated, w2,
                               expert_sizes_gpu, hidden_size);
 
     // Step 7: moe_combine_result

ex_engine/kernels/kernels.h

@@ -0,0 +1 @@
+../xllm_kernels/kernels.h

ex_engine/kernels/ops_api.h

@@ -0,0 +1 @@
+../xllm_kernels/ops_api.h

ex_engine/kernels/param.h

@@ -0,0 +1 @@
+../xllm_kernels/param.h

ex_engine/xllm_kernels/kernels.h

@@ -0,0 +1,11 @@
+/* Auto-generated aggregation header for xllm::kernel namespace.
+ * Equivalent to CMake cc_library(NAME kernels HDRS param.h ops_api.h).
+ * 
+ * AST Layer 3: kernel dispatch interface
+ * Called by: xllm_layers/ (Layer 2)
+ * Calls:     xllm_kernels/ilu/ (Layer 4)
+ */
+#pragma once
+
+#include "param.h"
+#include "ops_api.h"

ex_engine/xllm_kernels/ops_api.h

@@ -0,0 +1,177 @@
+/* Copyright 2025 The xLLM Authors. 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
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+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 "param.h"
+
+namespace xllm::kernel {
+
+static const std::string kActModeSilu = "silu";
+static const std::string kActModeGelu = "gelu";
+static const std::string kActModeQuickGelu = "quick_gelu";
+static const std::string kActModeSwish = "swish";
+
+void apply_rotary(RotaryParams& params);
+
+void active(ActivationParams& params);
+
+void reshape_paged_cache(ReshapePagedCacheParams& params);
+
+void reshape_from_cache(ReshapeFromCacheParams& params);
+
+// Quantize and store KV cache to paged cache (INT8 quantization)
+// Only supported on MLU backend
+void quant_to_paged_cache(ReshapePagedCacheParams& params);
+
+// Dequantize KV cache from paged cache (INT8 to FP16/BF16)
+// Only supported on MLU backend
+void dequant_from_paged_cache(ReshapeFromCacheParams& params);
+
+void fused_layernorm(FusedLayerNormParams& params);
+
+torch::Tensor matmul(MatmulParams& params);
+
+torch::Tensor group_gemm(GroupGemmParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
+    MoeFusedTopkParams& params);
+
+std::vector<torch::Tensor> moe_gen_idx(MoeGenIdxParams& params);
+
+torch::Tensor moe_expand_input(MoeExpandInputParams& params);
+
+torch::Tensor moe_combine_result(MoeCombineResultParams& params);
+
+torch::Tensor moe_all2all_gen_send_layout(
+    MoeAll2AllGenSendLayoutParams& params);
+
+std::vector<torch::Tensor> moe_all2all_gen_gather_index(
+    MoeAll2AllGenGatherIndexParams& params);
+
+std::vector<torch::Tensor> moe_all2all_create(MoeAll2AllCreateParams& params);
+
+void moe_all2all_init(MoeAll2AllInitParams& params);
+
+void moe_all2all_dispatch(MoeAll2AllDispatchParams& params);
+
+void moe_all2all_combine(MoeAll2AllCombineParams& params);
+
+void moe_all2all_destroy(MoeAll2AllDestroyParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor> scaled_quantize(
+    ScaledQuantizeParams& params);
+
+torch::Tensor scaled_matmul(ScaledMatmulParams& params);
+
+torch::Tensor apply_top_k_top_p(TopKPParams& params);
+
+torch::Tensor random_sample(RandomSampleParams& params);
+
+torch::Tensor rejection_sample(RejectionSampleParams& params);
+
+void masked_indexer_select_paged_kv(MaskedIndexerSelectPagedKVParams& params);
+
+void gather_split(GatherSplitParams& params);
+
+void fused_mla_q(FusedMlaQParams& params);
+
+void fused_mla_kv(FusedMlaKVParams& params);
+
+void fused_indexer_q(FusedIndexerQParams& params);
+
+void fused_indexer_k(FusedIndexerKParams& params);
+
+// L2 normalization along the last dimension
+torch::Tensor l2_norm(torch::Tensor& x, double eps = 1e-6);
+
+// TODO: NPU moe_init_routing_v2 is equivalent to moe_gen_idx + moe_expand_input
+// (and token_count/cusum outputs) on other backends.
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+moe_init_routing_v2(MoeInitRoutingV2Params& params);
+
+// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format
+// Returns: (quantized_output, scale)
+std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
+    Fp8ScaledQuantizeParams& params);
+
+// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels
+// Performs: c = (a @ b.T) with scales applied
+torch::Tensor fp8_scaled_matmul(Fp8ScaledMatmulParams& params);
+
+// Static scaled FP8 quantization helper
+// Quantizes input tensor to FP8 using a pre-computed scale factor
+void static_scaled_fp8_quant(StaticScaledFp8QuantParams& params);
+
+// Fused RMSNorm + Static FP8 Quantization
+// These fused operations combine RMSNorm and FP8 quantization to reduce memory
+// bandwidth by avoiding the intermediate write-back to global memory.
+
+// Fused RMSNorm + Static FP8 Quantization
+// Returns: FP8 quantized output tensor
+torch::Tensor rms_norm_static_fp8_quant(RmsNormStaticFp8QuantParams& params);
+
+// Fused Add + RMSNorm + Static FP8 Quantization (with residual)
+// Returns: tuple of (FP8 quantized output, updated residual)
+std::tuple<torch::Tensor, torch::Tensor> fused_add_rms_norm_static_fp8_quant(
+    FusedAddRmsNormStaticFp8QuantParams& params);
+
+std::pair<torch::Tensor, torch::Tensor> fused_gdn_gating(
+    FusedGdnGatingParams& params);
+
+std::pair<torch::Tensor, torch::Tensor> fused_recurrent_gated_delta_rule(
+    FusedRecurrentGatedDeltaRuleParams& params);
+
+torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params);
+
+torch::Tensor gated_layer_norm(GatedLayerNormParams& params);
+
+std::pair<torch::Tensor, torch::Tensor> partial_rotary_embedding(
+    PartialRotaryEmbeddingParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+fused_qkvzba_split_reshape_cat(FusedQkvzbaSplitReshapeParams& params);
+
+void gemma_rms_norm(GemmaRMSNormParams& params);
+
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
+split_qkv_rmsnorm_mrope(SplitQkvRmsnormMropeParams& params);
+
+bool has_split_qkv_rmsnorm_mrope_specialization(int64_t num_q_heads,
+                                                int64_t num_kv_heads,
+                                                int64_t head_size);
+
+torch::Tensor build_split_qkv_rmsnorm_mrope_gather_pattern(
+    int64_t rope_dim,
+    const std::vector<int64_t>& mrope_section,
+    bool is_interleaved,
+    const torch::Device& device);
+
+std::pair<torch::Tensor, torch::Tensor> chunk_gated_delta_rule(
+    ChunkGatedDeltaRuleParams& params);
+
+torch::Tensor recurrent_gated_delta_rule(
+    const torch::Tensor& query,
+    const torch::Tensor& key,
+    const torch::Tensor& value,
+    torch::Tensor& state,
+    const std::optional<torch::Tensor>& beta,
+    const std::optional<double> scale,
+    const std::optional<torch::Tensor>& actual_seq_lengths,
+    const std::optional<torch::Tensor>& ssm_state_indices,
+    const std::optional<torch::Tensor>& num_accepted_tokens,
+    const std::optional<torch::Tensor>& g,
+    const std::optional<torch::Tensor>& gk);
+}  // namespace xllm::kernel