Merge branch 'main' of https://github.com/dylanyunlon/project_6

This reverts commit 5c03156978.

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

qwen3_6_scripts/api_server.py

@@ -903,39 +903,6 @@ def build_app(args: Namespace) -> FastAPI:
         allow_headers=args.allowed_headers,
     )
 
-    @app.middleware("http")
-    async def sanitize_chat_body(request: Request, call_next):
-        """Strip fields from chat messages that vLLM's pydantic models reject.
-
-        Some replay datasets include ``index`` on messages (used by OpenAI
-        streaming deltas but forbidden by the non-streaming request schema).
-        Stripping it here avoids a ValidatorIterator 400 before our handler
-        even runs.
-        """
-        if (request.method == "POST"
-                and request.url.path.endswith("/v1/chat/completions")):
-            content_type = request.headers.get("content-type", "")
-            if "json" in content_type or not content_type:
-                try:
-                    body = await request.json()
-                    changed = False
-                    for msg in body.get("messages", []) if isinstance(body, dict) else []:
-                        if isinstance(msg, dict) and "index" in msg:
-                            del msg["index"]
-                            changed = True
-                    if changed:
-                        import json as _json
-                        raw = _json.dumps(body).encode("utf-8")
-
-                        async def patched_body():
-                            return raw
-
-                        request._body = raw
-                        request._receive = patched_body  # noqa
-                except Exception:
-                    pass
-        return await call_next(request)
-
     @app.exception_handler(RequestValidationError)
     async def validation_exception_handler(raw_request, exc):
         _bi100_log_request_validation_4xx(raw_request, exc)

qwen3_6_scripts/paged_attn.py

@@ -23,61 +23,6 @@ try:
 except ImportError:
     _corex_fused_paged_prefill = None
 
-# ---------------------------------------------------------------------------
-# Tier 0 prefill: ixformer native flash_attn_varlen_func
-# Sub 168 (competitor) uses this via corex_fa2.py:333 — single fused kernel
-# instead of our multi-tile Python loop. This is the #1 prefill bottleneck.
-# ---------------------------------------------------------------------------
-_ixformer_flash_attn_varlen = None
-_ixformer_flash_attn_kvcache = None
-_ixformer_paged_attn_v1 = None
-_ixformer_flash_attn_func = None
-try:
-    from ixformer.contrib.vllm_flash_attn import (
-        flash_attn_varlen_func as _ixformer_flash_attn_varlen,
-    )
-except (ImportError, AttributeError):
-    pass
-try:
-    from ixformer.contrib.vllm_flash_attn import (
-        flash_attn_with_kvcache as _ixformer_flash_attn_kvcache,
-    )
-except (ImportError, AttributeError):
-    pass
-try:
-    import ixformer.functions as _ixf_F
-    _ixformer_paged_attn_v1 = _ixf_F.vllm_single_query_cached_kv_attention
-except (ImportError, AttributeError):
-    pass
-# Tier 0.5: ixformer top-level flash_attn_func (non-varlen)
-# Probe confirmed: flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None,
-#                                   causal=False, return_attn_probs=False)
-# Available at ixformer.functions.flash_attn_func on BI-V100 real machine.
-# Not varlen — requires [batch, seqlen, nheads, headdim] layout.
-# For single-sequence prefill (competition concurrency=1), this replaces
-# the entire Python Q-tiling loop with one C++ kernel.
-try:
-    _ixformer_flash_attn_func = _ixf_F.flash_attn_func
-except (NameError, AttributeError):
-    try:
-        import ixformer.functions as _ixf_F2
-        _ixformer_flash_attn_func = _ixf_F2.flash_attn_func
-    except (ImportError, AttributeError):
-        pass
-
-# Tier 0.6: corex_fa2 dispatch (3-mode: packed prefill, paged decode, chunked)
-# This module wraps ix_bridge C++ and ixformer Python backends with proper
-# fallback chain. Import lazily — if corex_fa2 is not deployed, fall through.
-_corex_fa2_dispatch = None
-try:
-    from ex_engine.python.corex_fa2 import CoreXFA2 as _CoreXFA2Class
-    # Instantiate later when we know num_heads/head_dim
-except ImportError:
-    _CoreXFA2Class = None
-
-_USE_IXFORMER_FLASH_PREFILL = env_bool("BI100_USE_IXFORMER_FLASH_PREFILL", True)
-_LOGGED_IXFORMER_PREFILL = set()
-
 # from vllm.attention.ops.prefix_prefill import context_attention_fwd
 # NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT
 # imported here.  On Iluvatar BI-V100 that kernel hangs the GPU card
@@ -1800,169 +1745,6 @@ class PagedAttention:
             k_scale=k_scale,
             v_scale=v_scale,
         )
-        # -----------------------------------------------------------------
-        # Tier 0: ixformer flash_attn_varlen_func (cu_seqlens packed)
-        # This is what sub 168 uses via corex_fa2.py:333.
-        # Handles variable-length sequences in a single fused kernel.
-        # -----------------------------------------------------------------
-        if (_USE_IXFORMER_FLASH_PREFILL
-                and _ixformer_flash_attn_varlen is not None
-                and alibi_slopes is None
-                and sliding_window is None
-                and k_scale == 1.0 and v_scale == 1.0
-                and kv_cache_dtype == "auto"):
-            try:
-                batch_size = seq_lens_tensor.shape[0]
-                num_q_heads = query.shape[1]
-                head_dim = query.shape[2]
-                scale = head_dim ** -0.5
-
-                # Build cu_seqlens for packed varlen interface
-                # For prefill, all tokens are fresh — cu_seqlens covers full seq
-                q_lens = (query_start_loc[1:] - query_start_loc[:-1])
-                cu_seqlens_q = torch.zeros(
-                    batch_size + 1, dtype=torch.int32, device=query.device)
-                cu_seqlens_q[1:] = torch.cumsum(q_lens, dim=0).to(torch.int32)
-
-                # For context_lens=0 (pure prefill), k_seqlens == q_seqlens
-                # For context_lens>0 (chunked prefill), we need to handle
-                # the cached KV — but flash_attn_varlen handles only the
-                # fresh Q/K/V, not the paged cache. Fall through for that case.
-                all_zero_context = bool(context_lens.max().item() == 0)
-                if all_zero_context:
-                    max_seqlen = int(q_lens.max().item())
-                    output = _ixformer_flash_attn_varlen(
-                        q=query, k=key, v=value,
-                        cu_seqlens_q=cu_seqlens_q,
-                        cu_seqlens_k=cu_seqlens_q,
-                        max_seqlen_q=max_seqlen,
-                        max_seqlen_k=max_seqlen,
-                        softmax_scale=scale,
-                        causal=True)
-                    if "varlen_prefill" not in _LOGGED_IXFORMER_PREFILL:
-                        _LOGGED_IXFORMER_PREFILL.add("varlen_prefill")
-                        import logging
-                        logging.getLogger(__name__).info(
-                            "[BI100 PREFILL] ixformer flash_attn_varlen: "
-                            "B=%d Hq=%d D=%d max_q=%d — FUSED kernel active",
-                            batch_size, num_q_heads, head_dim, max_seqlen)
-                    return output
-            except Exception as _e:
-                if "varlen_error" not in _LOGGED_IXFORMER_PREFILL:
-                    _LOGGED_IXFORMER_PREFILL.add("varlen_error")
-                    import logging
-                    logging.getLogger(__name__).warning(
-                        "[BI100 PREFILL] ixformer flash_attn_varlen failed: "
-                        "%s — falling through to Tier 0.5", _e)
-
-        # -----------------------------------------------------------------
-        # Tier 0.5: ixformer flash_attn_func (non-varlen, batch layout)
-        # Probe confirmed available: flash_attn_func(q, k, v, ...)
-        # For single-sequence (batch=1) prefill, reshape to [1, seqlen, h, d]
-        # and call one C++ kernel. This replaces the entire Python Q-tiling
-        # loop which iterates hundreds of times for long prompts.
-        # -----------------------------------------------------------------
-        if (_USE_IXFORMER_FLASH_PREFILL
-                and _ixformer_flash_attn_func is not None
-                and alibi_slopes is None
-                and sliding_window is None
-                and k_scale == 1.0 and v_scale == 1.0
-                and kv_cache_dtype == "auto"):
-            try:
-                batch_size = seq_lens_tensor.shape[0]
-                num_q_heads = query.shape[1]
-                num_kv_heads = key.shape[1] if key.dim() == 3 else query.shape[1]
-                head_dim = query.shape[2]
-                scale = head_dim ** -0.5
-
-                all_zero_context = bool(context_lens.max().item() == 0)
-                if all_zero_context and batch_size == 1:
-                    # Single sequence, pure prefill — reshape to batch format
-                    total_q = query.shape[0]
-                    # flash_attn_func expects [batch, seqlen, nheads, headdim]
-                    q_4d = query.unsqueeze(0)  # [1, total_q, num_q_heads, head_dim]
-                    k_4d = key.unsqueeze(0)
-                    v_4d = value.unsqueeze(0)
-
-                    out_4d = _ixformer_flash_attn_func(
-                        q_4d, k_4d, v_4d,
-                        dropout_p=0.0,
-                        softmax_scale=scale,
-                        causal=True)
-                    output = out_4d.squeeze(0)  # [total_q, num_q_heads, head_dim]
-
-                    if "func_prefill" not in _LOGGED_IXFORMER_PREFILL:
-                        _LOGGED_IXFORMER_PREFILL.add("func_prefill")
-                        import logging
-                        logging.getLogger(__name__).info(
-                            "[BI100 PREFILL] ixformer flash_attn_func: "
-                            "B=1 Hq=%d D=%d seqlen=%d — FUSED kernel active",
-                            num_q_heads, head_dim, total_q)
-                    return output
-            except Exception as _e:
-                if "func_error" not in _LOGGED_IXFORMER_PREFILL:
-                    _LOGGED_IXFORMER_PREFILL.add("func_error")
-                    import logging
-                    logging.getLogger(__name__).warning(
-                        "[BI100 PREFILL] ixformer flash_attn_func failed: "
-                        "%s — falling through to Python Q-tiling", _e)
-
-        # -----------------------------------------------------------------
-        # Tier 1: corex_fa2 dispatch (3-mode: packed, paged decode, chunked)
-        # This wraps ix_bridge C++ and ixformer Python backends.
-        # -----------------------------------------------------------------
-        if (_USE_IXFORMER_FLASH_PREFILL
-                and _CoreXFA2Class is not None
-                and alibi_slopes is None
-                and sliding_window is None
-                and k_scale == 1.0 and v_scale == 1.0
-                and kv_cache_dtype == "auto"):
-            try:
-                batch_size = seq_lens_tensor.shape[0]
-                num_q_heads = query.shape[1]
-                num_kv_heads = key.shape[1] if key.dim() == 3 else num_q_heads
-                head_dim = query.shape[2]
-
-                all_zero_context = bool(context_lens.max().item() == 0)
-                if all_zero_context:
-                    q_lens = (query_start_loc[1:] - query_start_loc[:-1])
-                    cu_seqlens_q = torch.zeros(
-                        batch_size + 1, dtype=torch.int32,
-                        device=query.device)
-                    cu_seqlens_q[1:] = torch.cumsum(
-                        q_lens, dim=0).to(torch.int32)
-                    max_seqlen = int(q_lens.max().item())
-
-                    fa2 = _CoreXFA2Class(num_q_heads, num_kv_heads, head_dim)
-                    if fa2.is_available:
-                        output = fa2.packed_prefill(
-                            query, key, value,
-                            cu_seqlens_q, cu_seqlens_q,
-                            max_seqlen, max_seqlen,
-                            causal=True)
-                        if "corex_fa2" not in _LOGGED_IXFORMER_PREFILL:
-                            _LOGGED_IXFORMER_PREFILL.add("corex_fa2")
-                            import logging
-                            logging.getLogger(__name__).info(
-                                "[BI100 PREFILL] CoreXFA2 packed_prefill: "
-                                "B=%d Hq=%d Hkv=%d D=%d max_q=%d",
-                                batch_size, num_q_heads, num_kv_heads,
-                                head_dim, max_seqlen)
-                        return output
-            except Exception as _e:
-                if "corex_fa2_error" not in _LOGGED_IXFORMER_PREFILL:
-                    _LOGGED_IXFORMER_PREFILL.add("corex_fa2_error")
-                    import logging
-                    logging.getLogger(__name__).warning(
-                        "[BI100 PREFILL] CoreXFA2 failed: %s — "
-                        "falling through to Python Q-tiling", _e)
-
-        # -----------------------------------------------------------------
-        # Tier 2 (fallback): Python Q-tiling with online softmax
-        # This is the current default — functional but slow for long prompts.
-        # 107K prompt = ~400 tile iterations in Python, each launching
-        # multiple CUDA kernels. Sub 694 shows 190s TTFT for such requests.
-        # -----------------------------------------------------------------
         return PagedAttention._forward_prefix_pytorch(
             query, key, value,
             key_cache, value_cache,

qwen3_6_scripts/serving_chat.py

@@ -118,17 +118,12 @@ def _sequential_greedy_fanout_count(
     request: ChatCompletionRequest,
     max_num_seqs: int,
 ) -> int:
-    """Return the supported fan-out width, or zero.
-
-    When max_num_seqs=1 (competition fixed config), vLLM cannot schedule
-    n>1 natively.  We sequentially execute n independent n=1 requests and
-    merge them.  This works for any temperature — deterministic (temp=0)
-    produces identical choices, stochastic produces diverse ones.
-    """
+    """Return the supported deterministic fan-out width, or zero."""
     n = request.n if request.n is not None else 1
     if (
         max_num_seqs == 1
-        and 2 <= n <= 4
+        and n == 2
+        and request.temperature == 0
         and not request.stream
         and not request.use_beam_search
         and request.best_of is None
@@ -143,8 +138,8 @@ def _merge_sequential_chat_responses(
     request_id: str,
     created_time: int,
 ) -> ChatCompletionResponse:
-    if len(responses) < 2:
-        raise ValueError("fan-out requires at least two responses")
+    if len(responses) != 2:
+        raise ValueError("deterministic fan-out requires exactly two responses")
 
     first = responses[0]
     if any(response.model != first.model for response in responses):
@@ -402,16 +397,8 @@ class OpenAIServingChat(OpenAIServing):
             # OpenAI API: max_completion_tokens takes precedence over max_tokens
             if request.max_completion_tokens is not None and request.max_tokens is None:
                 request.max_tokens = request.max_completion_tokens
-            prompt_len = len(prompt_inputs["prompt_token_ids"])
-            default_max_tokens = self.max_model_len - prompt_len
-            # Clamp max_tokens so prompt + completion <= max_model_len.
-            # Without this, evaluation systems (e.g. OpenCompass) that send
-            # max_tokens=131072 get 400 errors when prompt+max_tokens exceeds
-            # max_model_len, resulting in 0 score on all academic benchmarks.
-            if default_max_tokens < 1:
-                default_max_tokens = 1
-            if request.max_tokens is not None and request.max_tokens > default_max_tokens:
-                request.max_tokens = default_max_tokens
+            default_max_tokens = self.max_model_len - len(
+                prompt_inputs["prompt_token_ids"])
             if request.use_beam_search:
                 sampling_params = request.to_beam_search_params(
                     default_max_tokens)
@@ -505,7 +492,7 @@ class OpenAIServingChat(OpenAIServing):
                 logger.error(
                     "Sequential greedy fan-out unexpectedly returned a stream")
                 return self.create_error_response(
-                    f"Failed to aggregate n={fanout_count} completion")
+                    "Failed to aggregate deterministic n=2 completion")
             responses.append(child_response)
 
         try:
@@ -516,11 +503,11 @@ class OpenAIServingChat(OpenAIServing):
             )
         except ValueError as error:
             logger.error(
-                "Sequential fan-out aggregation failed: %s",
+                "Sequential greedy fan-out aggregation failed: %s",
                 type(error).__name__,
             )
             return self.create_error_response(
-                f"Failed to aggregate n={fanout_count} completion")
+                "Failed to aggregate deterministic n=2 completion")
 
         if raw_request is not None:
             metadata = RequestResponseMetadata(