dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Merge branch 'main' of https://dev.modelhub.org.cn/dylanyunlong/project_6 into main

Browse Source
This commit is contained in:
root
2026-08-17 04:19:31 +00:00
parent cdcf115037 cb926707af
commit 512f384a49
10 changed files with 2743 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
ex_engine/kernels/kernels.h Symbolic link
View File

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

1
ex_engine/kernels/ops_api.h Symbolic link
View File

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

1
ex_engine/kernels/param.h Symbolic link
View File

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

11
ex_engine/xllm_kernels/kernels.h Normal file
View File

@@ -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"

1101
ex_engine/xllm_kernels/ops_api.cpp Normal file
View File

File diff suppressed because it is too large Load Diff

177
ex_engine/xllm_kernels/ops_api.h Normal file
View File

@@ -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

1441
ex_engine/xllm_kernels/param.h Normal file
View File

File diff suppressed because it is too large Load Diff

33
qwen3_6_scripts/api_server.py
View File

@@ -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)

218
qwen3_6_scripts/paged_attn.py
View File

@@ -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,

33
qwen3_6_scripts/serving_chat.py
View File

@@ -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(