Skip to content

[MODEL][PD] Integrate Nanbeige4.2 with PD serving - #77

Merged
bjf-frz merged 6 commits into
InferMatrixLab:mainfrom
kunkunblueberry:kunkun/nanbeige-pd
Oct 14, 2026
Merged

bjf-frz merged 6 commits into
InferMatrixLab:mainfrom
kunkunblueberry:kunkun/nanbeige-pd

Conversation

@kunkunblueberry

@kunkunblueberry kunkunblueberry commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Design

The model implementation, LLMEngine, P/D worker execution, and NIXL transport are unchanged. Both workers and the single-engine reference load the same local checkpoint. Tests use normal exit handling, without forced traces or loop overrides.

Validation

Linux, 2×RTX 4080 SUPER 32GB, Python 3.12.3, torch 2.7.0+cu128, Triton 3.3.0, NIXL 1.4.1/UCX:

  • python -m pytest tests/test_pd.py --run-gpu -k 'nanbeige and gpu' -v: 6 passed with generated tiny weights; 6 passed in official-checkpoint mode (five use official weights, EOS uses controlled tiny weights).
  • Official checkpoint: Nanbeige/Nanbeige4.2-3B, revision b82e54bd609793562a75cbf9337970a93369eab5; official file sizes and LFS SHA256 checked.
  • Verified exact greedy/seeded output agreement, single-output handoff, depth 2, EOS/length termination, cancellation at four lifecycle boundaries, immediate request-ID reuse, credit/block cleanup, and worker exit.
  • Existing two-GPU Ouro PD tests: 8 passed.
  • Full CPU regression before the final test-only additions: 411 passed, 29 skipped. Latest targeted CPU tests: 21 passed, 3 skipped. The final full Linux rerun could not be confirmed after the server disconnected.
  • Ruff lint/format and diff whitespace checks passed. No performance benchmarks were run.

Limitations

Qualification covers single-host 1P1D, BF16/Triton, LAST_EXITED KV, async execution without CUDA Graphs. HTTP serving, arbitrary mid-transfer fault injection, and Nanbeige-specific worker-failure/timeout cases remain unverified. Multi-worker, Huginn, speculative PD, and #64 architecture refactoring are outside scope.

Related

Refs #69

- Replace hardcoded OuroConfig reconstruction with model-agnostic approach
- Store only essential config fields (vocab_size, total_ut_steps, max_position_embeddings)
- Enables PD coordinator to work with any recurrent model (Ouro, Nanbeige, etc.)
- Preserves existing fingerprint validation for P/D worker consistency
- Add test comparing 1P1D Nanbeige execution against single-engine baseline
- Verify token ID matching, finish state, and resource cleanup
- Test uses fixed 2-loop configuration with greedy sampling
- Requires 2 GPUs and NIXL for execution
- Update model support section to include Nanbeige4.2-3B
- Document 1P1D configuration requirements and limitations
- Add test validation command for GPU+NIXL environments
- Clarify Ouro vs Nanbeige feature coverage
@kunkunblueberry
kunkunblueberry marked this pull request as ready for review September 30, 2026 16:48
@kunkunblueberry

kunkunblueberry commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor Author

Self-Review: Nanbeige4.2 × PD Integration

Scope

Architecture

  • Removes Ouro-specific OuroConfig hardcoding in PDEngine
  • Uses model-agnostic config extraction (vocab_size, total_ut_steps, max_position_embeddings)
  • Preserves fingerprint validation for P/D worker consistency
  • Minimal changes: 8 insertions, 2 deletions in engine.py

Behavior

  • No changes to existing Ouro PD behavior
  • PDEngine now accepts any recurrent model config
  • Existing PD tests remain unaffected

API / Contract

  • No public API changes
  • Worker-coordinator protocol unchanged
  • Backward compatible with Ouro

Tests

Verified on A100 server:

  • test_nanbeige_config_validation - PASSED
  • test_pd_configuration_validation (4 variants) - ALL PASSED
  • New test test_pd_nanbeige_1p1d_matches_single_engine - collected successfully
  • Code imports and logic verification - PASSED
  • Uses trace mode with explicit depth=2 for Nanbeige fixed 2-loop execution

Limitation:

  • End-to-end PD test requires NIXL (GPU-to-GPU KV transfer library)
  • Test framework is complete and ready for NIXL environment

Diff Review

  • No debug code, temporary logging, or commented code
  • No unrelated imports or file changes
  • No formatting-only changes
  • Clean 4-commit history with conventional format

Documentation

  • Updated docs/README.md with Nanbeige PD configuration
  • Documented 1P1D requirements, limitations, and test command
  • Clear about what's NOT supported (multi-worker, adaptive exit in PD)

Files Changed

vllm_rlt/pd/engine.py    |  8 insertions, 2 deletions
tests/test_pd.py          | 81 insertions, 6 deletions
docs/README.md            | 15 insertions, 2 deletions
---
Total: 3 files, 104(+) 10(-)

Ready for Review

  • Implementation complete
  • Tests validated (CPU + config level)
  • Documentation updated
  • Clean diff, no scope creep
  • Full end-to-end validation pending NIXL environment

Key Review Points:

  1. Line 152-160 in vllm_rlt/pd/engine.py - model-agnostic config extraction
  2. Line 344-424 in tests/test_pd.py - 1P1D correctness test with trace mode
  3. Line 25-41 in docs/README.md - Nanbeige PD documentation

Commits:

  • d6ec23f - refactor(pd): remove Ouro-specific config assumption
  • 4097a6e - test(pd): add Nanbeige4.2 1P1D correctness test
  • bc88782 - docs: add Nanbeige4.2 PD support documentation
  • 1dea57c - fix(test): correct trace mode and scope issues

…sues

- Replace invalid ExitConfig 'fixed' mode with 'trace' mode
- Provide explicit depth=2 trace for Nanbeige fixed 2-loop execution
- Add required trace_id parameter to add_request calls
- Fix stats variable scope (define before use)
- Trace format: [prefill_depth, decode_token1, decode_token2, ...]
Comment thread vllm_rlt/pd/engine.py Outdated
model_info = next(iter(self.peers.values())).info["model"]
self.model = SimpleNamespace(
config=OuroConfig(**next(iter(self.peers.values())).info["model"])
config=SimpleNamespace(

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.

Please replace the nested SimpleNamespace objects with an explicit, typed dataclass such as PDModelInfo containing the required model metadata. The coordinator should store this metadata directly instead of constructing a synthetic model.config object.

@bjf-frz

bjf-frz commented Oct 10, 2026

Copy link
Copy Markdown
Collaborator

Please also fix the Ruff and whitespace check failures, and the model-loading issue that prevents the new PD test from running. Add and run a real Nanbeige 1P1D E2E test with GPU + NIXL, verifying token agreement with the single-engine baseline, exit depths, cancellation, request-ID reuse, and resource cleanup.

Store typed coordinator metadata instead of a synthetic model config. Load matching local checkpoints through the existing P/D worker path and cover output parity, loop depths, cancellation, ID reuse, EOS, and cleanup.

Restore the existing prefix cleanup assertion and document the validated configuration and remaining coverage boundaries.
@kunkunblueberry

kunkunblueberry commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor Author

Thanks @bjf-frz for the careful review. Sorry I did not inspect the initial implementation and tests thoroughly enough, and overstated readiness based on configuration tests and collection alone. This update supersedes my earlier self-review.

Self-review — 829588b

  • Architecture: replaced nested SimpleNamespace model/config objects with typed PDModelInfo, stored directly by the coordinator. Nanbeige reuses the existing Ouro P/D workers and NIXL handoff; no model, LLMEngine, worker, or transport redesign.
  • Fixes: the E2E tests now load matching local checkpoints through the actual worker-loading contract, use normal exit handling rather than forced traces, restore the old prefix cleanup assertion, and pass Ruff/whitespace checks.
  • Correctness: real GPU + NIXL runs passed: 6 tiny-checkpoint tests, 5 official-checkpoint tests + 1 controlled EOS test, and 8 existing Ouro PD tests. Coverage includes token/finish agreement, depth 2, first-token handoff, seeded sampling, cancellation, immediate ID reuse, and resource/process cleanup.
  • CPU: latest targeted tests: 21 passed, 3 skipped. The prior full Linux run passed 411 tests.

The reported code/test issues are addressed in this commit. Arbitrary mid-transfer fault injection, Nanbeige worker-failure/timeout cases, and HTTP serving are not claimed as validated. No performance claims or unrelated refactoring. The PR description and recipe now reflect the actual evidence.

@bjf-frz

bjf-frz commented Oct 13, 2026

Copy link
Copy Markdown
Collaborator

resolve the conflict, please

bjf-frz commented Oct 14, 2026 •

Copy link
Copy Markdown
Collaborator

Performance measurements on PR head cedaedf (2026-10-14):

With this configuration, 1P1D substantially reduces decode latency at higher offered loads, but the single prefill worker becomes a bottleneck, increasing TTFT and reducing aggregate throughput.

Setup: official Nanbeige4.2-3B weights, 2×B300/NVLink, BF16/Triton, fixed two-loop execution, async scheduling, CUDA Graphs and prefix caching disabled. Compared two independent single-GPU replicas against 1P1D on the same PR commit, with identical per-worker KV capacity and sequence/token budgets. Requests alternate between 128 and 4096 input tokens and generate exactly 64 tokens using greedy decoding with EOS ignored. Each load uses Poisson arrivals over 60 seconds, repeated three times.

Values below are two replicas → 1P1D, reported as the median of the three per-run statistics:

Offered load Output throughput (tok/s) p95 TTFT p95 TPOT
0.25 requests/s 16.0 → 16.0 (unchanged) 3.05 → 3.08 s 23.5 → 24.7 ms (+5%)
0.5 requests/s 33.9 → 31.8 (−6%) 3.08 → 7.11 s 73.4 → 28.1 ms (−62%)
1 request/s* 53.1 → 40.9 (−23%) 4.24 → 26.35 s 226.9 → 29.7 ms (−87%)
  • At 0.25 requests/s, throughput is arrival-limited: equal throughput does not imply equal maximum capacity. There is no observed p95 decode-latency benefit at this load.
  • Throughput includes draining outstanding requests after arrivals stop. TTFT starts at actual engine submission; HTTP overhead is excluded. TPOT = (last-token time − first-token time)/63. At 1 request/s, the replica baseline has submission lag; including that lag gives p95 arrival-to-first-token latency of 4.55 → 26.35 s, preserving the conclusion. Tail percentiles are descriptive for these samples.
  • Single-request TTFT is broadly similar. E2E latency for 128/1024/4096 input tokens is 0.764/0.950/4.532 s → 0.800/1.004/4.543 s. These small differences also include run-period and GPU-role differences and cannot be attributed entirely to PD.
  • Prefill capacity is the main bottleneck: both replica GPUs can perform prefill, whereas 1P1D dedicates only one GPU to it. A single 4096-token request spends about 3.043 s in the P stage. Separate ops-only/full profiles show attention accounts for about 90% of summed GPU kernel time on P. P-side queuing and handoff progress worsen under load. For example, the same 53 requests finish in 63.90 s with replicas versus 82.95 s with PD, explaining the throughput loss. Profiled timings were not used for performance comparisons.

*Correctness limitation: the official-checkpoint PD/single-engine agreement test passed, and single-request and 0.25/0.5 requests/s outputs match the reference. At 1 request/s, 6/166 replica-baseline requests diverge from the single-request greedy reference starting at output token 33. The cause remains unresolved, so this load point is diagnostic rather than a completed correctness qualification. No inference implementation changes were made.

These findings apply to the tested Triton, long-input/short-output configuration. FA4, CUDA Graphs, and longer outputs remain untested. The next priority is to validate a faster prefill backend, then repeat the PD comparison with the same GPU budget.

@bjf-frz

bjf-frz commented Oct 14, 2026

Copy link
Copy Markdown
Collaborator

The performance benefit looks limited: lower decode latency comes at the cost of higher TTFT and lower overall throughput under the tested workload.

@kunkunblueberry

Copy link
Copy Markdown
Contributor Author

Thank you @bjf-frz for taking the time to run these measurements and profiles. I'm very sorry for the delay in completing proper performance validation on my side; my access to suitable multi-GPU resources has been limited, and I should have communicated the validation gap more clearly instead of presenting the PR as ready too early.

I agree with your conclusion: under the tested configuration, lower decode latency does not compensate for the higher TTFT and reduced aggregate throughput. This PR should be described as Nanbeige integration with the existing PD runtime, not as a demonstrated overall performance improvement. The unresolved high-load output divergence also means that point cannot be treated as correctness-qualified.

I appreciate the detailed evidence and the suggested next step of validating a faster prefill backend before repeating the matched comparison. Sorry again for leaving this additional validation work to you, and thank you for your patience and careful review.

@bjf-frz

bjf-frz commented Oct 14, 2026

Copy link
Copy Markdown
Collaborator

lgtm

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