feat: add embedded-python model executor for qwen3 on cuda.#1891
feat: add embedded-python model executor for qwen3 on cuda.#1891zhang-minchao wants to merge 34 commits into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces a Python model executor framework that allows executing Python-defined model graphs, such as Qwen3, within xLLM's C++ worker using an embedded CPython interpreter. It adds the C++ bridge, a lightweight config reflection system, and the Python-side model, layers, op dispatching, and graph runner. Feedback on the changes focuses on strict adherence to the repository's style guide, including using project-root-relative paths for #include directives, preferring fixed-width integers and emplace_back in C++, using CHECK instead of TORCH_CHECK, ensuring complete Python type annotations, and routing Python diagnostic logs through the shared logger instead of print().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
7430c11 to
479ab8a
Compare
c962bc8 to
eb22c0b
Compare
be4d1a1 to
357910e
Compare
| "task", | ||
| "device_id", | ||
| "devices", | ||
| "python_model_path", |
There was a problem hiding this comment.
why define python_model_path?
There was a problem hiding this comment.
python_model_path is optional. The server embeds CPython and starts the model from the C++ binary, so source-tree and custom deployments cannot assume the xLLM python package is already on sys.path. This setting lets the embedded interpreter prepend the directory containing that package before import; when it is empty, we fall back to XLLM_PYTHON_MODEL_PATH and then the existing Python import path. Keeping it in ModelConfig also makes CLI/JSON configuration and config dumps consistent.
Rename the KV-cache memory knob and change its meaning to the fraction of GPU memory left FREE after model weights are loaded (previously a whole-GPU utilization cap that reserved total*(1-util) as buffer). This makes the KV budget independent of weight size and matches the new semantics across the LLM / VLM / OneRec / LLM-Rec engines and the xtensor phy-page allocator. - keep --max_memory_utilization as a deprecated alias (gflag + pybind arg) that overrides the new flag under the new semantics - propagate the rename through options, c_api/cc_api, and config parsing - update config_json_test to the new key
Add a Python model-execution path selectable via --model_impl=python --python_model_path=<dir>: the C++ PyCausalLM embeds CPython and runs an nn.Module graph (xllm_models package) whose every op dispatches to the same fused xllm_ops CUDA kernels the native C++ decoder uses, so outputs stay byte-identical to the native path. - C++ bridge: PyCausalLM + py_model_bridge (weight sharding via the model's sharding_plan, paged-KV/attention state exposed through a thread-local forward context), registered in model_registry behind --model_impl - fused ops exposed as torch.ops.xllm_ops.* (rms_norm, fused_add_rms_norm, fused_qk_norm_rope, silu_and_mul, attention, all_reduce/all_gather); mutating kernels use honest in-place (a!) void schemas (no wrapper clone) - Python package: Qwen3 model/layers, ops dispatch, register_fake impls, tensor-parallel collectives, and an optional XLLM_TC_BACKEND switch for torch.compile (cudagraphs / inductor / reduce-overhead) plus an optional XLLM_USE_TRITON path for silu_and_mul - add xllm_ops_test unit target
PyCausalLM::forward unconditionally rebuilt AttentionMetadata via AttentionMetadataBuilder::build, discarding the executor-supplied parameters.attn_metadata. In --enable_graph mode CudaGraphExecutorImpl pre-builds a persistent, bucket-padded and flashinfer-planned metadata; the rebuilt one has plan_info == nullptr, so attention flashinfer-plans during capture and fatals with 'operation not permitted when stream is capturing'. Reuse parameters.attn_metadata when present (mirrors the native qwen3 reuse-if-present pattern), building our own only on the eager path where it is null. Verified on cuda: all 6 decode buckets (num_tokens 1..32) capture cleanly and greedy output stays byte-identical to the native C++ graph path (MD5 993c223fce58).
…nput lengths
Enabling the python-model prefill piecewise graph and running the full
py-vs-cpp regression (input 512/2k/8k x tp 1/2 x graph on/off) exposed three
shared C++/flashinfer crashes in graph-mode prefill (cpp native and python
model crash identically -> not python-introduced). All three are root-caused
and fixed here. They are interdependent (fixing only a subset still crashes),
so they land together.
1. fa3/SM90 ragged prefill crashed under piecewise graph replay
("TMA descriptor 700 / BatchPrefillWithRaggedKVCacheSM90Run illegal memory
access"). Graph-mode prefill was the only path still on fa3 (decode/chunked
already pin fa2):
- cuda_graph_executor_impl.cpp: graph-mode prefill backend -> fa2
- batch_prefill.cpp: derive run backend from the plan's uri (_sm90 => fa3
else fa2) so run follows plan instead of re-deriving independently
- flashinfer_planinfo.cpp: branch plan() arg-count on backend (fa3=16,
fa2=19), not hardware is_support_sm90a(), fixing fa2-on-SM90
"Expected 19 but got 16 arguments"
2. tp>=2 large-bucket piecewise replay SIGSEGV in cuGraphLaunch. The custom
SharedVMMAllocator graph pool is a monotonic bump allocator (no-op free)
that reuses physical memory across shapes via virtual-space switching --
sound for decode's single-graph-per-shape capture, unsafe for piecewise's
29 interleaved sub-graphs under one shape. Piecewise now uses torch's
standard private graph pool (graph_pool_); decode keeps the VMM pool.
3. fa2 graph-prefill at 8192 tokens overflowed the 128MB flashinfer float
workspace ("Buffer overflow ... batch_prefill_tmp_s"). Raised default
flashinfer_workspace_buffer_size 128MB -> 512MB (fa2 split-k needs more
scratch than fa3 did).
NOTE: change (1) also switches cpp NATIVE graph-mode prefill from fa3 to fa2
(a production behavior change), justified because fa3 graph-prefill crashed.
Result: the full 24-cell Cartesian runs with zero crashes; tp=1 graph is
exact py-vs-cpp parity, tp=2 graph is within the conc=1 cross-card jitter floor.
Roll back the kv_cache_memory_fraction knob introduced on this branch and restore main's max_memory_utilization semantics across the config, options, engines, allocator, c_api/cc_api and pybind surfaces. This keeps the python-model-executor PR scoped to the executor and off unrelated KV-cache sizing behavior.
Roll back the 512MB flashinfer split-k workspace bump. The larger buffer was tied to graph-mode prefill sizing that does not belong in this PR; restore the 128MB default and its original description.
Rename the embedded-CPython model package from xllm_models to python and split
it into a layered dependency chain models -> {layers, model_runner} -> ops ->
kernels:
- layers/: RMSNorm, RotaryEmbedding, Column/RowParallelLinear, embedding
(split out of the former flat layers.py)
- ops/: torch.ops.xllm_ops dispatch (dispatch.py) + register_fake / attention
disallow_in_graph (fake_impls.py)
- kernels/: leaf JIT Triton kernels
- model_runner/: GraphRunner owning torch.compile / cudagraph capture
(extracted from qwen3.py), mirroring SGLang's model_executor runners
Update the C++ bridge import strings (py_causal_lm, py_model_bridge) and the
model_impl flag doc to the new package name, ship the package inside the wheel
under the xLLM install tree (setup.py copytree), and default
XLLM_PYTHON_MODEL_PATH to the launcher directory (launch_server.py).
Document how the python model executor registers xllm_ops across hardware (one device-agnostic schema + per-DispatchKey impl + shared register_fake, CUDA today, PrivateUse1 for Ascend NPU next) and why the Python model graph is fully reusable across backends. Grounded in the local vLLM (direct_register_custom_op / current_platform.dispatch_key) and SGLang (TORCH_LIBRARY_IMPL(npu, PrivateUse1)) implementations.
- use project-root-relative include paths (core/..., models/...) - CHECK (glog) instead of TORCH_CHECK; add missing glog include - fixed-width ints (int -> int32_t) in the sharding-plan executor - emplace_back over push_back - add type annotations on all python model/layer/op functions - drop the env-gated _instr tp-parity debug scaffolding (parity already validated; scripts.logger is not importable in the embedded runtime)
Replace the declarative sharding_plan (C++ interprets a data structure) with an imperative load_weights (Python drives all transforms directly). - Expose StateDict to Python via PYBIND11_EMBEDDED_MODULE(xllm_weight_loader) - PyCausalLM::load_model reduced to wrapping StateDict + calling load_weights - Delete ShardingPlan type, sharding_plan(), load_assembled_weights() - Qwen3 load_weights does TP slice, QKV/gate-up fuse, copy — all in Python This unifies weight transforms (slice, fuse, future NZ/quant) as plain torch ops in Python, making them extensible without C++ changes.
The copytree destination used os.path.dirname(cmake_dir) which resolves to build/, but the wheel packs from extdir (build/lib.linux-...-cpython-312/). Fix by using extdir directly so bdist_wheel includes xllm/python/*.py.
…cutor Replace the broken torch.compile(backend="cudagraphs") path with a manual decode full-graph capture (bucketed by batch size) and torch.compile piecewise prefill (graph-break at attention). This matches the C++ CudaGraphExecutorImpl design while keeping all capture/replay orchestration in Python. - Add --python_graph_backend / --python_graph_max_batch gflags to replace the unreliable XLLM_TC_BACKEND / XLLM_GRAPH_MAX_BATCH env vars; config is passed to Python via build_config_dict (no env var dependency). - Add bool enable_cuda_graph param to the update_decode_plan C++ op so flashinfer produces fixed-layout plans for graph replay. - Rewrite GraphRunner: "cudagraphs" mode routes decode to DecodeFullGraphRunner (manual torch.cuda.CUDAGraph per bucket) and prefill to torch.compile(dynamic=True) piecewise. - Two capture bugs fixed: (1) replay after capture (capture only records), (2) plan + capture + replay on one dedicated stream (ordering guarantee).
…prefill - Fix CudaGraphExecutorImpl illegal memory access when bucket > actual batch: pad per-sequence metadata (indptr, last_page_len, kv_seq_lens) for padding sequences so FlashInfer plan covers all padded tokens. - Size persistent tensors by bucketed max_seqs_per_batch. - Update llm_decode_metadata_update kernel to fill dummy metadata for padding sequences (indptr constant, last_page_len=1, delta=0). - Switch py_graph prefill from piecewise CUDA graph to eager forward (eliminates TTFT regression at high batch while keeping decode TPOT benefit from full graph). - Add @torch._dynamo.disable on PagedAttention.forward to prevent recompilation per layer_id. - Make --model_impl and --python_model_path visible in --help output.
…adding. - delete breakable_cuda_graph.py (piecewise capture infrastructure, unused) - remove PrefillPiecewiseRunner and related from graph_runner.py - inline attention dispatch into PagedAttention.forward(), drop @eager_break - pre-allocate pad_indptr_buf to eliminate per-step torch.arange allocation - clean up fake_impls.py docstring referencing removed piecewise semantics
…dules. - split dispatch.py into compute.py (4 compute ops), attention.py (7 attention ops), collectives.py (TP all_reduce/all_gather as torch custom ops) - register fake_impl and disallow_in_graph inline with each op binding - register all_reduce/all_gather as proper torch.library custom ops for Dynamo graph capture - route layers/attention.py through ops.* instead of direct torch.ops.xllm_ops.* - delete dispatch.py, fake_impls.py, forward_batch.py (dead code) - remove all XLLM_USE_TRITON environment variable dependencies
…s bridge - Delete xllm_attention_ops.cpp/h (batch_decode/prefill/update_plan torch ops were only used by the Python executor, now replaced by flashinfer Python API) - Move reshape_paged_cache registration into xllm_ops_library.cpp - layers/attention.py: use flashinfer BatchDecodeWithPagedKVCacheWrapper / BatchPrefillWithRaggedKVCacheWrapper directly instead of torch.ops.xllm_ops - Rename py_model_bridge → py_model_helper, consolidate PyDictVisitor and dtype_to_string utilities - Move py_causal_lm to models/llm/ (where other CausalLM implementations live) - Add core/kernels/xllm_torch_ops.h as device-agnostic dispatch header - Remove USE_CUDA guard from PyCausalLM compilation (device-agnostic) - Add FlashinferWorkspace skip for model_impl=python in llm_worker_impl.cpp
The backend == "fa3" change was unnecessary — main's hardware-based branching is correct for the ragged prefill path.
The uri-based backend detection was unnecessary — main's determine_attention_backend() call is correct.
Remove outdated references to deleted xllm_attention_ops.cpp, redundant argumentation, and premature NPU landing checklist. Keep only the current architecture table, minimal extension guide, and risk notes.
Add complete description of the three interaction phases (init config dict, weight loading, per-step forward args) and remove speculative PrivateUse1 references.
The variable holds attention metadata, not generic metadata. Rename for clarity and update the design doc to match.
Add model_impl field to ModelContext so model_registry.cpp no longer reads the global FLAGS_model_impl directly. The caller (llm_worker_impl) sets it before calling create_llm_model.
- initialize PyTorch NCCL TP collectives and honor JSON executor settings. - prevent KV padding overflow and remove obsolete graph batch configuration. - add config, dump, and embedded Python regression coverage.
Update the Python executor design document to match the current CUDA, TP, JSON, and decode graph implementation.
99a0f54 to
6f8e5e6
Compare
Install FlashInfer 0.6.14 after the AOT build and make the embedded Python runtime share the xLLM TVM-FFI libraries.
| fi | ||
|
|
||
| # Install tvm-ffi | ||
| RUN set -eux; \ |
There was a problem hiding this comment.
remove tvm-ffi installation manually.
| # Install flashinfer + AOT compile | ||
| ENV FLASHINFER_VERSION=${FLASHINFER_VERSION} | ||
| ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST} | ||
| RUN set -eux; \ |
- cpp_framework_python_model_architecture.md: document the architecture decision to add python model execution on the c++ serving framework and why python owns both the model and its ModelExecutor. - refactor_python_model_impl.md: describe the implementable refactor plan covering cross-language interface, object lifecycle, migration order and acceptance conditions, baselined on PR xLLM-AI#1891.
Description
Add an embedded-CPython model executor that runs Qwen3 inference on CUDA through the xLLM runtime, as an opt-in alternative to the native C++ model path. The C++ engine drives an in-process Python interpreter to execute the model forward pass, while all compute stays on the existing
xllm_opsfused CUDA kernels, so the Python path is numerically validated against the native path (byte-identical greedy decode at tp=1, and within the cross-card jitter floor at tp=2).Enable it with
--model_impl=python --python_model_path=<path-to-xllm>; the default native path is unchanged.Main pieces:
PyCausalLM+py_model_helper(pybind) let the C++ executor construct and call a Python model like any nativeCausalLM, sharing the KV cache, forward inputs, and attention metadata.xllm/python/) — Qwen3 model, layers (linear / layernorm / embedding / rotary), op dispatch, triton kernels, and a cudagraph-capablemodel_runner, all built onxllm_ops.PROPERTYstruct reflection (property_reflect.h+macros.h) — dependency-free field reflection behind an opt-inREFLECT_PROPERTIES, so the fullModelArgs/ParallelArgscross the pybind boundary without hand-maintained field lists.xllm_opsop registration — ops layer structured as compute/collectives modules; TP collectives registered astorch.library.custom_opfor Dynamo capture; attention uses flashinfer Python API directly.CudaGraphExecutorImplpadding crash (bucket > actual batch caused illegal memory access due to missing per-sequence metadata for padding sequences).服务启动
镜像构建
docker/Dockerfile.cuda已在现有 FlashInfer AOT 构建完成后安装 Python executor运行时。AOT 源码版本保持为
0.6.2,全局 Python package 固定为flashinfer-python==0.6.14。FlashInfer 的
apache-tvm-ffi依赖会在 Python package 内携带另一份libtvm_ffi.so,而 xLLM 二进制链接的是/usr/local/lib/libtvm_ffi.so。glibc 按 inode 区分动态库,两份不同 inode 的TVM-FFI 会重复执行全局注册并导致进程 abort。Dockerfile 现在会:
RECORD的旧packaging,再安装flashinfer-python==0.6.14;sysconfig动态定位 site-packages,不硬编码 Python 版本路径;libtvm_ffi_testing.so复制到/usr/local/lib;libtvm_ffi.so和libtvm_ffi_testing.so软链到/usr/local/lib中的同一份库;Performance
Cross-batch-size regression on Qwen3-0.6B, dedicated H100 (cuda:6), in=1024 out=256, CUDA_LAUNCH_BLOCKING=0. Four variants: cpp_eager, cpp_graph, py_eager, py_graph. Batch list deliberately mixes bucket boundaries and odd values to exercise the decode-graph padding path.
Takeaways:
plan()+cudaGraphLaunch(), GPU-bound, structurally identical.Related Issues
Change Type