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

fix(CRITICAL): engine death on image request + stop overwriting base corex modules

Browse Source 
Root cause from latest docker build log:
  ValueError: You set image=0 in --limit-mm-per-prompt, but found 1 items
  → Engine background task crashes → AsyncEngineDeadError → all subsequent 503

Fixes:
1. computility-run.yaml: add --limit-mm-per-prompt image=1
   Prevents multimodal ValueError from killing the engine process.

2. patch_ops.sh: DON'T overwrite base image's corex_gdn.py/corex_moe.py
   Comp 168 log proves base image's corex modules work with libcorex_gdn.so.
   Our overwrite broke CoreXGDN.__init__ (unexpected kwarg 'num_v_heads').
   Only deploy ours if base has NO corex modules at all.
   Also deploy corex_fa2.py if base lacks it.

3. qwen3_5.py: try multiple CoreXGDN init signatures
   Base image CoreXGDN may accept different kwargs than ours.
   Try kwargs form first, fall back to positional.

4. corex_gdn.py: accept both calling conventions in __init__
   Future-proof for when we DO need to deploy ours.

5. Copied upstream_ref headers: ilu_layer_fused_moe.h, ilu_layer_attention.h
   Last 2 missing ILU files from xllm. All 14/14 now present.
This commit is contained in:
Claude
2026-08-10 09:12:05 +00:00
parent ff3562b941
commit f87689a4ef
7 changed files with 435 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

82
ex_engine/include/ilu_layer_attention.h Normal file
View File

@@ -0,0 +1,82 @@
/* 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 <torch/torch.h>
#include <tuple>
#include "framework/kv_cache/kv_cache.h"
#include "framework/model/model_input_params.h"
#include "layers/common/attention_metadata.h"
namespace xllm {
namespace layer {
class AttentionImpl : public torch::nn::Module {
public:
AttentionImpl() = default;
AttentionImpl(int64_t num_heads,
int64_t head_size,
float scale,
int64_t num_kv_heads,
int64_t sliding_window);
AttentionImpl(int64_t num_heads,
int64_t head_size,
int64_t num_kv_heads,
int64_t v_head_dim,
int64_t sliding_window,
float scale,
bool use_fused_mla_qkv,
bool enable_lighting_indexer,
bool enable_mla);
std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
const AttentionMetadata& attn_metadata,
torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& value,
KVCache& kv_cache);
void prefill_forward(torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& value,
torch::Tensor& output,
const torch::Tensor& k_cache,
const std::optional<torch::Tensor>& v_cache,
const AttentionMetadata& attn_metadata);
void decoder_forward(torch::Tensor& query,
torch::Tensor& output,
const torch::Tensor& k_cache,
const std::optional<torch::Tensor>& v_cache,
const AttentionMetadata& attn_metadata);
private:
int64_t num_heads_;
int64_t head_size_;
float scale_;
int64_t num_kv_heads_;
int64_t v_head_dim_;
bool use_fused_mla_qkv_;
bool enable_lighting_indexer_;
bool enable_mla_;
int64_t sliding_window_;
};
TORCH_MODULE(Attention);
} // namespace layer
} // namespace xllm

131
ex_engine/include/ilu_layer_fused_moe.h Normal file
View File

@@ -0,0 +1,131 @@
/* 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 <torch/torch.h>
#include "framework/model/model_args.h"
#include "framework/model/model_input_params.h"
#include "framework/parallel_state/parallel_args.h"
#include "framework/quant_args.h"
#include "framework/state_dict/state_dict.h"
#include "framework/state_dict/utils.h"
#include "layers/common/deep_ep.h"
#include "layers/common/dense_mlp.h"
#include "layers/common/fused_moe_base.h"
#include "layers/common/linear.h"
#include "platform/device.h"
#include "util/tensor_helper.h"
namespace xllm {
namespace layer {
class FusedMoEImpl : public torch::nn::Module {
public:
FusedMoEImpl() = default;
FusedMoEImpl(const ModelArgs& model_args,
const FusedMoEArgs& moe_args,
const QuantArgs& quant_args,
const ParallelArgs& parallel_args,
const torch::TensorOptions& options);
torch::Tensor forward_experts(const torch::Tensor& hidden_states,
const torch::Tensor& router_logits,
bool enable_all2all_communication);
torch::Tensor forward(const torch::Tensor& hidden_states,
const ModelInputParams& input_params);
void load_state_dict(const StateDict& state_dict);
private:
// struct to store the selected expert info
struct SelectedExpertInfo {
torch::Tensor reduce_weight;
torch::Tensor combine_idx;
torch::Tensor token_count_slice;
std::optional<torch::Tensor> cusum_token_count;
std::optional<torch::Tensor> input_scale;
};
// initial steps for MoE computation, select the experts for each token
torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
const torch::Tensor& router_logits_2d,
SelectedExpertInfo& selected_expert_info,
bool enable_all2all_communication);
private:
int64_t num_total_experts_;
int64_t topk_;
int64_t num_expert_group_;
int64_t topk_group_;
double route_scale_;
int64_t hidden_size_;
int64_t n_shared_experts_;
bool is_gated_;
int64_t renormalize_;
std::string hidden_act_;
std::string scoring_func_;
bool is_smoothquant_;
int64_t num_experts_per_rank_;
int64_t start_expert_id_;
// Deep EP related parameters
bool enable_deep_ep_;
DeepEPBuffer deep_ep_buffer_;
DeepEPParams deep_ep_params_;
torch::Tensor dispatch_recv_token_tensor_head_;
torch::Tensor dispatch_recv_token_tensor_tail_;
// steams for parallel shared experts
std::unique_ptr<Stream> shared_stream_;
std::unique_ptr<Stream> routed_stream_;
xllm::Device device_;
bool stream_initialized_ = false;
ReplicatedLinear gate_{nullptr};
DenseMLP shared_experts_{nullptr};
DeepEP deep_ep_{nullptr};
QuantArgs quant_args_;
ParallelArgs parallel_args_;
torch::TensorOptions options_;
ProcessGroup* tp_pg_;
DEFINE_WEIGHT(w13);
DEFINE_FUSED_WEIGHT(w1);
DEFINE_FUSED_WEIGHT(w3);
DEFINE_FUSED_WEIGHT(w2);
DEFINE_WEIGHT(e_score_correction_bias);
DEFINE_WEIGHT(w13_scale);
DEFINE_FUSED_WEIGHT(w1_scale);
DEFINE_FUSED_WEIGHT(w3_scale);
DEFINE_FUSED_WEIGHT(w2_scale);
DEFINE_FUSED_WEIGHT(input_smooth);
DEFINE_FUSED_WEIGHT(act_smooth);
void load_e_score_correction_bias(const StateDict& state_dict);
void load_experts(const StateDict& state_dict);
// create the group gemm output tensor with the workspace
torch::Tensor create_group_gemm_output(const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& group_list,
torch::ScalarType dtype,
torch::Tensor& workspace);
};
TORCH_MODULE(FusedMoE);
} // namespace layer
} // namespace xllm

93
ex_engine/python/corex_gdn.py
View File

@@ -82,18 +82,33 @@ class CoreXGDN:
def __init__(
self,
num_heads: int,
head_dim: int,
num_heads: int = 0,
head_dim: int = 128,
layer_idx: int = 0,
chunk_size: int = 16,
eps: float = 1e-6,
# kwargs from qwen3_5.py (GatedDeltaNet uses separate k/v dims)
num_v_heads: int = 0,
num_k_heads: int = 0,
head_k_dim: int = 0,
head_v_dim: int = 0,
conv_kernel_size: int = 4,
**kwargs, # future-proof
):
self.num_heads = num_heads
self.head_dim = head_dim
# Accept both calling conventions:
# CoreXGDN(num_heads, head_dim) — simple
# CoreXGDN(num_v_heads=.., num_k_heads=.., head_k_dim=.., head_v_dim=..) — from qwen3_5.py
self.num_v_heads = num_v_heads or num_heads
self.num_k_heads = num_k_heads or num_heads
self.head_k_dim = head_k_dim or head_dim
self.head_v_dim = head_v_dim or head_dim
self.num_heads = self.num_v_heads
self.head_dim = self.head_k_dim
self.layer_idx = layer_idx
self.chunk_size = chunk_size
self.eps = eps
self.scale = head_dim ** -0.5
self.conv_kernel_size = conv_kernel_size
self.scale = self.head_k_dim ** -0.5
self._decode_warned = False
self._prefill_warned = False
@@ -106,20 +121,68 @@ class CoreXGDN:
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
gate: torch.Tensor,
beta: torch.Tensor,
conv_state: Optional[torch.Tensor],
temporal_state: Optional[torch.Tensor],
hidden_states: torch.Tensor,
attn_metadata,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
conv_state: torch.Tensor,
temporal_state: torch.Tensor,
in_proj_qkv, # nn.Module — projects hidden → conv_dim
in_proj_z, # nn.Module — projects hidden → val_dim
in_proj_b, # nn.Module — projects hidden → num_v_heads (beta)
in_proj_a, # nn.Module — projects hidden → num_v_heads (alpha/dt)
conv1d_weight, # (conv_dim, 1, kernel_size) depthwise conv weight
A_log, # (num_v_heads,) log decay parameters
dt_bias, # (num_v_heads,) dt bias
norm, # GatedRMSNorm module
out_proj, # RowParallelLinear
) -> torch.Tensor:
"""Full GDN layer forward — matches qwen3_5.py calling convention.
This mirrors the PyTorch _pytorch_forward() path but uses ixformer
matmul acceleration and fused CoreX GDN ops when available.
"""
from vllm.model_executor.parallel_utils.communication_op import (
tensor_model_parallel_all_reduce,
)
try:
from vllm.distributed import get_tensor_model_parallel_world_size
except ImportError:
get_tensor_model_parallel_world_size = lambda: 1
tp_size = get_tensor_model_parallel_world_size()
local_key_dim = self.num_k_heads * self.head_k_dim // tp_size
local_val_dim = self.num_v_heads * self.head_v_dim // tp_size
local_num_v = self.num_v_heads
local_num_k = self.num_k_heads
local_conv_dim = local_key_dim * 2 + local_val_dim
is_prefill = getattr(attn_metadata, 'num_prefill_tokens', 0) > 0
# Project all tokens at once
mixed_qkv_all, _ = in_proj_qkv(hidden_states)
z_all, _ = in_proj_z(hidden_states)
b_all, _ = in_proj_b(hidden_states)
a_all, _ = in_proj_a(hidden_states)
if is_prefill:
return self._prefill(q, k, v, gate, beta, temporal_state)
if not self._prefill_warned:
logger.info("Using fused CoreX GDN prefill operator")
self._prefill_warned = True
return self._full_prefill(
hidden_states, attn_metadata, conv_state, temporal_state,
mixed_qkv_all, z_all, b_all, a_all,
conv1d_weight, A_log, dt_bias, norm, out_proj,
local_key_dim, local_val_dim, local_num_v, local_num_k,
local_conv_dim)
else:
return self._decode(q, k, v, gate, beta, conv_state, temporal_state)
if not self._decode_warned:
logger.info("Using fused CoreX GDN decode operator")
self._decode_warned = True
return self._full_decode(
hidden_states, attn_metadata, conv_state, temporal_state,
mixed_qkv_all, z_all, b_all, a_all,
conv1d_weight, A_log, dt_bias, norm, out_proj,
local_key_dim, local_val_dim, local_num_v, local_num_k,
local_conv_dim)
def _prefill(self, q, k, v, gate, beta, temporal_state):
if not self._prefill_warned: