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

feat(CRITICAL): import wudixzy/competition complete corex stack — 12 prebuilt .so + 13 CUDA kernels + 2615-line qwen3_5.py

Browse Source 
Source: github.com/wudixzy/competition (1527 files, BI-V100 competition reference)

Imported assets:
- 12 prebuilt CoreX .so extensions (corex-3.2.3-ivcore10):
  corex_gdn_{beta_decay,causal_conv,gated_norm,packed_decode,qk_map}.so
  corex_moe_{direct_routed,exact_reduce,weight_gather}.so
  corex_attn_head_rms_norm.so, corex_paged_kv_gather.so
  corex_block_major_kv_transfer.so, corex_fused_paged_prefill.so

- 13 CUDA kernel sources (.cu) for above extensions
- 11 build scripts (build_corex_*.sh)
- install_prebuilt_corex.sh (SHA256-verified .so deployment)
- qwen3_5.py (2615 lines) with FULL corex kernel integration
- 9 vllm vendor override files (block manager, sampler, etc)
- 19 patch scripts (model_runner, xformers, block_major, etc)
- Complete serving layer (serving_chat, protocol, api_server, etc)
- bi100_env.py, bi100_profile.py, gdn_prefix.py, block_major_kv_cache.py
- Dockerfile aligned with reference build chain
- computility-run.yaml with BI100_MOE_COREX_DIRECT_ROUTED=1

Call chain verified:
  Dockerfile COPY → patch_ops.sh → install_prebuilt_corex.sh → 12 .so to $VLLM_ROOT
  qwen3_5.py imports: from vllm import corex_gdn_* / corex_moe_* / corex_attn_*
This commit is contained in:
project6-dev
2026-08-11 03:55:38 +00:00
parent 81875fff52
commit 5862708b32
86 changed files with 24702 additions and 9860 deletions
Download Patch File Download Diff File Expand all files Collapse all files

436
qwen3_6_scripts/patch_model_runner.py
View File

@@ -1,46 +1,43 @@
"""
Fix: prefix_cache_hit stays True for chunked-prefill chunk 2+ even when past cache.
"""Patch vLLM 0.6.3 prefix-cache and MRoPE chunk alignment bugs."""
Root cause:
model_runner.py _compute_for_prefix_cache_hit has three cases:
Case 1: prefix_cache_len <= context_len → "already past cache, do normal"
Case 2: context_len < prefix_cache_len < seq_len → partial hit, correct
Case 3: seq_len <= prefix_cache_len → full hit, reduce to 1 token
from __future__ import annotations
Case 1 does nothing (leaves prefix_cache_hit = True). Then in utils.py:
if inter_data.prefix_cache_hit:
block_table = computed_block_nums ← ONLY the original prefix blocks!
import pathlib
But context_len > prefix_cache_len means chunk 1 tokens (between prefix_cache_len
and context_len) are ALSO in KV cache and need to be in block_table.
block_table = computed_block_nums misses all chunk-1 blocks.
from patch_utils import package_root, replace_once
In _forward_prefix_pytorch:
num_ctx_blocks = ceil(context_len / block_size) # e.g. 268
block_tables.shape[1] = len(computed_block_nums) # e.g. 12 <-- too small!
At tile_blk >= 12: blk_ids is empty → k_t shape [..., 0] → amax crash.
Fix:
Set prefix_cache_hit = False for Case 1, so utils.py falls through to:
elif chunked_prefill_enabled:
block_table = block_tables[seq_id] ← full block table (prefix + chunk1)
"""
HELPER_ANCHOR = """\
logger = init_logger(__name__)
import re
import sys
LORA_WARMUP_RANK = 8"""
CANDIDATE_PATHS = [
"/usr/local/corex/lib64/python3/dist-packages/vllm/worker/model_runner.py",
"/usr/local/corex/lib/python3/dist-packages/vllm/worker/model_runner.py",
]
HELPER_REPLACEMENT = """\
logger = init_logger(__name__)
OLD_BLOCK = """\
def _slice_mrope_positions(positions, start, stop, expected_len):
if positions is None or len(positions) != 3:
raise RuntimeError("MRoPE positions must contain three axes")
sliced = [axis[start:stop] for axis in positions]
lengths = [len(axis) for axis in sliced]
if lengths != [expected_len] * 3:
raise RuntimeError(
"MRoPE/input token length mismatch after chunk alignment: "
f"positions={lengths}, input_tokens={expected_len}, "
f"slice=({start}, {stop})")
return sliced
LORA_WARMUP_RANK = 8"""
PREFIX_PAST_ANCHOR = """\
if prefix_cache_len <= context_len:
# We already passed the cache hit region,
# so do normal computation.
pass"""
NEW_BLOCK = """\
PREFIX_PAST_REPLACEMENT = """\
if prefix_cache_len <= context_len:
# We already passed the cache hit region,
# so do normal computation.
@@ -51,28 +48,361 @@ NEW_BLOCK = """\
# causing an empty blk_ids slice and a zero-dim amax() crash.
inter_data.prefix_cache_hit = False"""
import os
PARTIAL_HIT_ANCHOR = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][uncomputed_start:]
context_len = prefix_cache_len
patched = False
for path in CANDIDATE_PATHS:
if not os.path.exists(path):
continue
with open(path, "r") as f:
src = f.read()
if OLD_BLOCK not in src:
if NEW_BLOCK in src:
print(f"[patch_model_runner] already patched: {path}")
patched = True
break
print(f"[patch_model_runner] WARNING: expected block not found in {path}, skipping")
continue
patched_src = src.replace(OLD_BLOCK, NEW_BLOCK, 1)
with open(path, "w") as f:
f.write(patched_src)
print(f"[patch_model_runner] patched Case-1 prefix_cache_hit fix in: {path}")
patched = True
break
inter_data.context_lens[seq_idx] = context_len
inter_data.query_lens[
seq_idx] = inter_data.seq_lens[seq_idx] - context_len"""
if not patched:
print("[patch_model_runner] ERROR: could not find model_runner.py at any known path", file=sys.stderr)
sys.exit(1)
PARTIAL_HIT_REPLACEMENT = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][uncomputed_start:]
context_len = prefix_cache_len
inter_data.context_lens[seq_idx] = context_len
inter_data.query_lens[
seq_idx] = inter_data.seq_lens[seq_idx] - context_len
if inter_data.mrope_input_positions is not None:
positions = inter_data.mrope_input_positions[seq_idx]
if positions is not None:
inter_data.mrope_input_positions[seq_idx] = \\
_slice_mrope_positions(
positions, uncomputed_start, None,
inter_data.query_lens[seq_idx])"""
FULL_HIT_ANCHOR = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][-1:]
inter_data.query_lens[seq_idx] = 1
inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1"""
FULL_HIT_REPLACEMENT = """\
inter_data.input_positions[seq_idx] = inter_data.input_positions[
seq_idx][-1:]
inter_data.query_lens[seq_idx] = 1
inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1
if inter_data.mrope_input_positions is not None:
positions = inter_data.mrope_input_positions[seq_idx]
if positions is not None:
inter_data.mrope_input_positions[seq_idx] = \\
_slice_mrope_positions(positions, -1, None, 1)"""
MULTIMODAL_MROPE_ANCHOR = """\
mrope_input_positions, mrope_position_delta = \\
MRotaryEmbedding.get_input_positions(
token_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
image_token_id=hf_config.image_token_id,
video_token_id=hf_config.video_token_id,
vision_start_token_id=hf_config.vision_start_token_id,
vision_end_token_id=hf_config.vision_end_token_id,
spatial_merge_size=hf_config.vision_config.
spatial_merge_size,
context_len=inter_data.context_lens[seq_idx],
)
seq_data.mrope_position_delta = mrope_position_delta
inter_data.mrope_input_positions[
seq_idx] = mrope_input_positions"""
MULTIMODAL_MROPE_REPLACEMENT = """\
# vLLM 0.6.3 returns positions through the end of token_ids,
# while chunked prefill sends only [context_len:seq_len].
# Compute the full MRoPE map once so the delta remains tied to
# the complete request, then select exactly the physical query.
mrope_input_positions, mrope_position_delta = \\
MRotaryEmbedding.get_input_positions(
token_ids,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
image_token_id=hf_config.image_token_id,
video_token_id=hf_config.video_token_id,
vision_start_token_id=hf_config.vision_start_token_id,
vision_end_token_id=hf_config.vision_end_token_id,
spatial_merge_size=hf_config.vision_config.
spatial_merge_size,
context_len=0,
)
mrope_input_positions = _slice_mrope_positions(
mrope_input_positions,
inter_data.context_lens[seq_idx],
inter_data.seq_lens[seq_idx],
len(inter_data.input_tokens[seq_idx]))
seq_data.mrope_position_delta = mrope_position_delta
inter_data.mrope_input_positions[
seq_idx] = mrope_input_positions"""
MODEL_INPUT_FIELDS_ANCHOR = """\
multi_modal_kwargs: Optional[BatchedTensorInputs] = None
request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None"""
MODEL_INPUT_FIELDS_REPLACEMENT = """\
multi_modal_kwargs: Optional[BatchedTensorInputs] = None
# BI100 scheduler-owned GDN prefix-cache actions. These plain Python
# objects are included in the multiprocess model-input broadcast.
gdn_restore_key: Optional[Tuple[int, bytes]] = None
gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
gdn_segment_offsets: Optional[List[int]] = None
request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None"""
BASE_BROADCAST_ANCHOR = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
return tensor_dict
@classmethod"""
BASE_BROADCAST_REPLACEMENT = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"gdn_restore_key\": self.gdn_restore_key,
\"gdn_capture_points\": self.gdn_capture_points,
\"gdn_evict_keys\": self.gdn_evict_keys,
\"gdn_segment_offsets\": self.gdn_segment_offsets,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
return tensor_dict
@classmethod"""
SAMPLING_BROADCAST_ANCHOR = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
_add_sampling_metadata_broadcastable_dict(tensor_dict,
self.sampling_metadata)"""
SAMPLING_BROADCAST_REPLACEMENT = """\
\"multi_modal_kwargs\": self.multi_modal_kwargs,
\"gdn_restore_key\": self.gdn_restore_key,
\"gdn_capture_points\": self.gdn_capture_points,
\"gdn_evict_keys\": self.gdn_evict_keys,
\"gdn_segment_offsets\": self.gdn_segment_offsets,
\"prompt_adapter_mapping\": self.prompt_adapter_mapping,
\"prompt_adapter_requests\": self.prompt_adapter_requests,
\"virtual_engine\": self.virtual_engine,
\"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
\"finished_requests_ids\": self.finished_requests_ids,
}
_add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
_add_sampling_metadata_broadcastable_dict(tensor_dict,
self.sampling_metadata)"""
BUILDER_INIT_ANCHOR = """\
self.finished_requests_ids = finished_requests_ids
self.decode_only = True
# Intermediate data"""
BUILDER_INIT_REPLACEMENT = """\
self.finished_requests_ids = finished_requests_ids
self.decode_only = True
self.gdn_restore_key = None
self.gdn_capture_points = None
self.gdn_evict_keys = None
self.gdn_segment_offsets = None
# Intermediate data"""
ADD_SEQ_GROUP_ANCHOR = """\
def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
\"\"\"Add a sequence group to the builder.\"\"\"
seq_ids = seq_group_metadata.seq_data.keys()"""
ADD_SEQ_GROUP_REPLACEMENT = """\
def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
\"\"\"Add a sequence group to the builder.\"\"\"
gdn_actions = (
seq_group_metadata.gdn_restore_key,
seq_group_metadata.gdn_capture_points,
seq_group_metadata.gdn_evict_keys,
seq_group_metadata.gdn_segment_offsets,
)
if any(value is not None for value in gdn_actions):
if not seq_group_metadata.is_prompt:
raise RuntimeError(\"GDN prefix-cache actions require prefill\")
if any(value is not None for value in (
self.gdn_restore_key, self.gdn_capture_points,
self.gdn_evict_keys, self.gdn_segment_offsets)):
raise RuntimeError(
\"only one GDN prefix-cache action group is supported\")
(self.gdn_restore_key, self.gdn_capture_points,
self.gdn_evict_keys, self.gdn_segment_offsets) = gdn_actions
seq_ids = seq_group_metadata.seq_data.keys()"""
BUILD_RESULT_ANCHOR = """\
lora_mapping=lora_mapping,
lora_requests=lora_requests,
multi_modal_kwargs=multi_modal_kwargs,
request_ids_to_seq_ids=request_ids_to_seq_ids,"""
BUILD_RESULT_REPLACEMENT = """\
lora_mapping=lora_mapping,
lora_requests=lora_requests,
multi_modal_kwargs=multi_modal_kwargs,
gdn_restore_key=self.gdn_restore_key,
gdn_capture_points=self.gdn_capture_points,
gdn_evict_keys=self.gdn_evict_keys,
gdn_segment_offsets=self.gdn_segment_offsets,
request_ids_to_seq_ids=request_ids_to_seq_ids,"""
EXECUTE_KWARGS_ANCHOR = """\
seqlen_agnostic_kwargs = {
\"finished_requests_ids\": model_input.finished_requests_ids,
\"request_ids_to_seq_ids\": model_input.request_ids_to_seq_ids,
} if self.has_inner_state else {}
if (self.observability_config is not None"""
EXECUTE_KWARGS_REPLACEMENT = """\
seqlen_agnostic_kwargs = {
\"finished_requests_ids\": model_input.finished_requests_ids,
\"request_ids_to_seq_ids\": model_input.request_ids_to_seq_ids,
} if self.has_inner_state else {}
gdn_prefix_kwargs = {}
if model_input.gdn_restore_key is not None:
gdn_prefix_kwargs[\"gdn_restore_key\"] = model_input.gdn_restore_key
if model_input.gdn_capture_points is not None:
gdn_prefix_kwargs[\"gdn_capture_points\"] = (
model_input.gdn_capture_points)
if model_input.gdn_evict_keys is not None:
gdn_prefix_kwargs[\"gdn_evict_keys\"] = model_input.gdn_evict_keys
if model_input.gdn_segment_offsets is not None:
gdn_prefix_kwargs[\"gdn_segment_offsets\"] = (
model_input.gdn_segment_offsets)
if (self.observability_config is not None"""
MODEL_CALL_ANCHOR = """\
**MultiModalInputs.as_kwargs(multi_modal_kwargs,
device=self.device),
**seqlen_agnostic_kwargs)"""
MODEL_CALL_REPLACEMENT = """\
**MultiModalInputs.as_kwargs(multi_modal_kwargs,
device=self.device),
**seqlen_agnostic_kwargs,
**gdn_prefix_kwargs)"""
PROFILE_KV_LAYERS_ANCHOR = """\
num_layers = self.model_config.get_num_layers(self.parallel_config)"""
PROFILE_KV_LAYERS_REPLACEMENT = """\
num_layers = self.model_config.get_num_attention_layers(
self.parallel_config)"""
def patch_model_runner(model_runner: pathlib.Path) -> None:
replace_once(
model_runner,
HELPER_ANCHOR,
HELPER_REPLACEMENT,
required=True,
already_contains="def _slice_mrope_positions(",
)
replace_once(
model_runner,
PREFIX_PAST_ANCHOR,
PREFIX_PAST_REPLACEMENT,
required=True,
already_contains="Must clear prefix_cache_hit so _add_seq_group",
)
replace_once(
model_runner,
PARTIAL_HIT_ANCHOR,
PARTIAL_HIT_REPLACEMENT,
required=True,
already_contains="positions, uncomputed_start, None,",
)
replace_once(
model_runner,
FULL_HIT_ANCHOR,
FULL_HIT_REPLACEMENT,
required=True,
already_contains="_slice_mrope_positions(positions, -1, None, 1)",
)
replace_once(
model_runner,
MULTIMODAL_MROPE_ANCHOR,
MULTIMODAL_MROPE_REPLACEMENT,
required=True,
already_contains="Compute the full MRoPE map once",
)
replace_once(
model_runner,
MODEL_INPUT_FIELDS_ANCHOR,
MODEL_INPUT_FIELDS_REPLACEMENT,
already_contains="gdn_restore_key: Optional[Tuple[int, bytes]]",
)
replace_once(
model_runner,
BASE_BROADCAST_ANCHOR,
BASE_BROADCAST_REPLACEMENT,
already_contains=BASE_BROADCAST_REPLACEMENT,
)
replace_once(
model_runner,
SAMPLING_BROADCAST_ANCHOR,
SAMPLING_BROADCAST_REPLACEMENT,
already_contains=SAMPLING_BROADCAST_REPLACEMENT,
)
replace_once(
model_runner,
BUILDER_INIT_ANCHOR,
BUILDER_INIT_REPLACEMENT,
already_contains="self.gdn_restore_key = None",
)
replace_once(
model_runner,
ADD_SEQ_GROUP_ANCHOR,
ADD_SEQ_GROUP_REPLACEMENT,
already_contains="gdn_actions = (",
)
replace_once(
model_runner,
BUILD_RESULT_ANCHOR,
BUILD_RESULT_REPLACEMENT,
already_contains="gdn_restore_key=self.gdn_restore_key",
)
replace_once(
model_runner,
EXECUTE_KWARGS_ANCHOR,
EXECUTE_KWARGS_REPLACEMENT,
already_contains="gdn_prefix_kwargs = {}",
)
replace_once(
model_runner,
MODEL_CALL_ANCHOR,
MODEL_CALL_REPLACEMENT,
already_contains="**gdn_prefix_kwargs)",
)
replace_once(
model_runner,
PROFILE_KV_LAYERS_ANCHOR,
PROFILE_KV_LAYERS_REPLACEMENT,
required=True,
already_contains=PROFILE_KV_LAYERS_REPLACEMENT,
)
if __name__ == "__main__":
patch_model_runner(package_root("vllm") / "worker" / "model_runner.py")