Skip to content

feat: add embedded-python model executor for qwen3 on cuda.#1891

Open
zhang-minchao wants to merge 34 commits into
xLLM-AI:mainfrom
zhang-minchao:feat/python-model-executor
Open

feat: add embedded-python model executor for qwen3 on cuda.#1891
zhang-minchao wants to merge 34 commits into
xLLM-AI:mainfrom
zhang-minchao:feat/python-model-executor

Conversation

@zhang-minchao

@zhang-minchao zhang-minchao commented Jul 6, 2026

Copy link
Copy Markdown
Collaborator

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_ops fused 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:

  • C++ ↔ Python bridgePyCausalLM + py_model_helper (pybind) let the C++ executor construct and call a Python model like any native CausalLM, sharing the KV cache, forward inputs, and attention metadata.
  • Python model package (xllm/python/) — Qwen3 model, layers (linear / layernorm / embedding / rotary), op dispatch, triton kernels, and a cudagraph-capable model_runner, all built on xllm_ops.
  • PROPERTY struct reflection (property_reflect.h + macros.h) — dependency-free field reflection behind an opt-in REFLECT_PROPERTIES, so the full ModelArgs / ParallelArgs cross the pybind boundary without hand-maintained field lists.
  • xllm_ops op registration — ops layer structured as compute/collectives modules; TP collectives registered as torch.library.custom_op for Dynamo capture; attention uses flashinfer Python API directly.
  • Graph-mode decode + prefill fixes — decode uses manual full CUDA graph (bucketed by batch, with proper padding metadata for FlashInfer); prefill runs eager for minimal TTFT. Also fixes the native C++ CudaGraphExecutorImpl padding crash (bucket > actual batch caused illegal memory access due to missing per-sequence metadata for padding sequences).

服务启动

# Python 模型执行器 (eager)
./xllm --model /path/to/Qwen3-0.6B \
  --model_impl=python \
  --python_model_path=/path/to/xllm/xllm \
  --device_id=0 \
  --enable_graph=false

# Python 模型执行器 (decode full CUDA graph)
./xllm --model /path/to/Qwen3-0.6B \
  --model_impl=python \
  --python_model_path=/path/to/xllm/xllm \
  --device_id=0 \
  --enable_graph=false \
  --python_graph_backend=cudagraphs

镜像构建

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 现在会:

  • 覆盖基础镜像中无 pip RECORD 的旧 packaging,再安装
    flashinfer-python==0.6.14
  • 通过 sysconfig 动态定位 site-packages,不硬编码 Python 版本路径;
  • libtvm_ffi_testing.so 复制到 /usr/local/lib
  • 将 Python package 内的 libtvm_ffi.solibtvm_ffi_testing.so 软链到
    /usr/local/lib 中的同一份库;
  • 在镜像构建阶段通过 device/inode 比较确认两条路径解析到同一文件。

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.

batch bucket pad? cpp_eager TTFT cpp_eager TPOT cpp_graph TTFT cpp_graph TPOT py_eager TTFT py_eager TPOT py_graph TTFT py_graph TPOT py_graph vs cpp_graph TTFT py_graph vs cpp_graph TPOT
1 1 6.158 2.735 6.241 1.994 8.166 5.069 8.237 2.013 +32.0% +0.9%
2 2 8.946 2.814 8.446 1.935 8.680 5.172 9.643 2.038 +14.2% +5.4%
3 4 y 11.430 2.723 11.903 2.151 11.982 5.166 11.915 2.221 +0.1% +3.2%
5 8 y 15.639 2.844 14.311 2.076 14.945 5.246 14.418 2.137 +0.7% +3.0%
7 8 y 22.455 2.890 25.782 2.197 22.665 5.282 18.812 2.267 -27.0% +3.2%
8 8 19.253 2.949 24.809 2.211 26.060 5.281 25.338 2.232 +2.1% +0.9%
13 16 y 27.455 2.961 23.962 2.564 31.308 5.424 29.876 2.614 +24.7% +2.0%
16 16 29.456 3.123 31.349 2.607 28.833 5.505 33.355 2.646 +6.4% +1.5%
31 32 y 53.353 3.613 48.408 3.502 48.089 5.799 57.553 3.531 +18.9% +0.8%
47 48 y 70.293 4.485 64.080 4.431 62.662 6.055 62.854 4.458 -1.9% +0.6%
100 112 y 128.417 7.380 123.170 7.224 125.745 7.991 130.316 7.396 +5.8% +2.4%

Takeaways:

  • py_graph TPOT ≈ cpp_graph TPOT (within +1~5%) — both decode paths replay a captured CUDA graph; the hot path is plan() + cudaGraphLaunch(), GPU-bound, structurally identical.
  • py_graph TPOT vs py_eager TPOT: −55% to −60% at low batch (1→8) — decode full graph eliminates per-step CPython/host dispatch overhead.
  • py_graph TTFT ≈ py_eager TTFT — prefill runs eager (no piecewise graph overhead).
  • cpp_graph fully working — previously crashed with "illegal memory access" at non-bucket-boundary batch sizes; fixed by padding per-sequence FlashInfer metadata for dummy sequences.
  • Padding rows (pad?=y) verify that decode graph correctly handles padded dummy sequences (slot_mapping=0 to reserved block, output sliced to actual batch).

Related Issues

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread xllm/core/common/macros.h Outdated
Comment thread tests/core/kernels/cuda/xllm_ops_test.cpp Outdated
Comment thread xllm/core/kernels/cuda/xllm_ops_library.cpp Outdated
Comment thread xllm/core/kernels/cuda/xllm_ops_library.cpp Outdated
Comment thread xllm/models/py_causal_lm.cpp Outdated
Comment thread xllm/python/ops/fake_impls.py Outdated
Comment thread xllm/python/ops/fake_impls.py Outdated
Comment thread xllm/python/ops/fake_impls.py Outdated
Comment thread xllm/python/ops/fake_impls.py Outdated
Comment thread xllm/python/registry.py Outdated
@zhang-minchao zhang-minchao force-pushed the feat/python-model-executor branch from 7430c11 to 479ab8a Compare July 6, 2026 08:16
@zhang-minchao zhang-minchao changed the title Feat/python model executor feat: add embedded-python model executor for qwen3 on cuda. Jul 6, 2026
@zhang-minchao zhang-minchao force-pushed the feat/python-model-executor branch 2 times, most recently from c962bc8 to eb22c0b Compare July 6, 2026 11:25
@zhang-minchao zhang-minchao marked this pull request as draft July 6, 2026 12:22
@zhang-minchao zhang-minchao marked this pull request as ready for review July 7, 2026 19:18
@zhang-minchao zhang-minchao force-pushed the feat/python-model-executor branch from be4d1a1 to 357910e Compare July 7, 2026 19:19
"task",
"device_id",
"devices",
"python_model_path",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why define python_model_path?

@zhang-minchao zhang-minchao Jul 10, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread xllm/core/layers/cuda/flashinfer_planinfo.cpp Outdated
Comment thread xllm/models/model_registry.cpp Outdated
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.
Comment thread xllm/core/kernels/cuda/cuda_ops_library.cpp
- 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.
@zhang-minchao zhang-minchao force-pushed the feat/python-model-executor branch from 99a0f54 to 6f8e5e6 Compare July 10, 2026 07:23
Install FlashInfer 0.6.14 after the AOT build and make the embedded Python runtime share the xLLM TVM-FFI libraries.
Comment thread docker/Dockerfile.cuda
fi

# Install tvm-ffi
RUN set -eux; \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

remove tvm-ffi installation manually.

Comment thread docker/Dockerfile.cuda
# Install flashinfer + AOT compile
ENV FLASHINFER_VERSION=${FLASHINFER_VERSION}
ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST}
RUN set -eux; \

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto.

- 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants