Merge branch 'main' of https://dev.modelhub.org.cn/dylanyunlong/project_6 into main
This commit is contained in:
1
ex_engine/kernels/kernels.h
Symbolic link
1
ex_engine/kernels/kernels.h
Symbolic link
@@ -0,0 +1 @@
|
|||||||
|
../xllm_kernels/kernels.h
|
||||||
1
ex_engine/kernels/ops_api.h
Symbolic link
1
ex_engine/kernels/ops_api.h
Symbolic link
@@ -0,0 +1 @@
|
|||||||
|
../xllm_kernels/ops_api.h
|
||||||
1
ex_engine/kernels/param.h
Symbolic link
1
ex_engine/kernels/param.h
Symbolic link
@@ -0,0 +1 @@
|
|||||||
|
../xllm_kernels/param.h
|
||||||
11
ex_engine/xllm_kernels/kernels.h
Normal file
11
ex_engine/xllm_kernels/kernels.h
Normal 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
1101
ex_engine/xllm_kernels/ops_api.cpp
Normal file
File diff suppressed because it is too large
Load Diff
177
ex_engine/xllm_kernels/ops_api.h
Normal file
177
ex_engine/xllm_kernels/ops_api.h
Normal 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
1441
ex_engine/xllm_kernels/param.h
Normal file
File diff suppressed because it is too large
Load Diff
@@ -903,39 +903,6 @@ def build_app(args: Namespace) -> FastAPI:
|
|||||||
allow_headers=args.allowed_headers,
|
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)
|
@app.exception_handler(RequestValidationError)
|
||||||
async def validation_exception_handler(raw_request, exc):
|
async def validation_exception_handler(raw_request, exc):
|
||||||
_bi100_log_request_validation_4xx(raw_request, exc)
|
_bi100_log_request_validation_4xx(raw_request, exc)
|
||||||
|
|||||||
@@ -23,61 +23,6 @@ try:
|
|||||||
except ImportError:
|
except ImportError:
|
||||||
_corex_fused_paged_prefill = None
|
_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
|
# from vllm.attention.ops.prefix_prefill import context_attention_fwd
|
||||||
# NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT
|
# NOTE: context_attention_fwd (Triton kernel from prefix_prefill.py) is NOT
|
||||||
# imported here. On Iluvatar BI-V100 that kernel hangs the GPU card
|
# imported here. On Iluvatar BI-V100 that kernel hangs the GPU card
|
||||||
@@ -1800,169 +1745,6 @@ class PagedAttention:
|
|||||||
k_scale=k_scale,
|
k_scale=k_scale,
|
||||||
v_scale=v_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(
|
return PagedAttention._forward_prefix_pytorch(
|
||||||
query, key, value,
|
query, key, value,
|
||||||
key_cache, value_cache,
|
key_cache, value_cache,
|
||||||
|
|||||||
@@ -118,17 +118,12 @@ def _sequential_greedy_fanout_count(
|
|||||||
request: ChatCompletionRequest,
|
request: ChatCompletionRequest,
|
||||||
max_num_seqs: int,
|
max_num_seqs: int,
|
||||||
) -> int:
|
) -> int:
|
||||||
"""Return the supported fan-out width, or zero.
|
"""Return the supported deterministic 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.
|
|
||||||
"""
|
|
||||||
n = request.n if request.n is not None else 1
|
n = request.n if request.n is not None else 1
|
||||||
if (
|
if (
|
||||||
max_num_seqs == 1
|
max_num_seqs == 1
|
||||||
and 2 <= n <= 4
|
and n == 2
|
||||||
|
and request.temperature == 0
|
||||||
and not request.stream
|
and not request.stream
|
||||||
and not request.use_beam_search
|
and not request.use_beam_search
|
||||||
and request.best_of is None
|
and request.best_of is None
|
||||||
@@ -143,8 +138,8 @@ def _merge_sequential_chat_responses(
|
|||||||
request_id: str,
|
request_id: str,
|
||||||
created_time: int,
|
created_time: int,
|
||||||
) -> ChatCompletionResponse:
|
) -> ChatCompletionResponse:
|
||||||
if len(responses) < 2:
|
if len(responses) != 2:
|
||||||
raise ValueError("fan-out requires at least two responses")
|
raise ValueError("deterministic fan-out requires exactly two responses")
|
||||||
|
|
||||||
first = responses[0]
|
first = responses[0]
|
||||||
if any(response.model != first.model for response in responses):
|
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
|
# OpenAI API: max_completion_tokens takes precedence over max_tokens
|
||||||
if request.max_completion_tokens is not None and request.max_tokens is None:
|
if request.max_completion_tokens is not None and request.max_tokens is None:
|
||||||
request.max_tokens = request.max_completion_tokens
|
request.max_tokens = request.max_completion_tokens
|
||||||
prompt_len = len(prompt_inputs["prompt_token_ids"])
|
default_max_tokens = self.max_model_len - len(
|
||||||
default_max_tokens = self.max_model_len - prompt_len
|
prompt_inputs["prompt_token_ids"])
|
||||||
# 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
|
|
||||||
if request.use_beam_search:
|
if request.use_beam_search:
|
||||||
sampling_params = request.to_beam_search_params(
|
sampling_params = request.to_beam_search_params(
|
||||||
default_max_tokens)
|
default_max_tokens)
|
||||||
@@ -505,7 +492,7 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
logger.error(
|
logger.error(
|
||||||
"Sequential greedy fan-out unexpectedly returned a stream")
|
"Sequential greedy fan-out unexpectedly returned a stream")
|
||||||
return self.create_error_response(
|
return self.create_error_response(
|
||||||
f"Failed to aggregate n={fanout_count} completion")
|
"Failed to aggregate deterministic n=2 completion")
|
||||||
responses.append(child_response)
|
responses.append(child_response)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -516,11 +503,11 @@ class OpenAIServingChat(OpenAIServing):
|
|||||||
)
|
)
|
||||||
except ValueError as error:
|
except ValueError as error:
|
||||||
logger.error(
|
logger.error(
|
||||||
"Sequential fan-out aggregation failed: %s",
|
"Sequential greedy fan-out aggregation failed: %s",
|
||||||
type(error).__name__,
|
type(error).__name__,
|
||||||
)
|
)
|
||||||
return self.create_error_response(
|
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:
|
if raw_request is not None:
|
||||||
metadata = RequestResponseMetadata(
|
metadata = RequestResponseMetadata(
|
||||||
|
|||||||
Reference in New Issue
Block a user