Skip to content

[REFACTOR]M7 3/3: validate sampling and RNG contracts on CUDA - #153

Draft
ycjcx123 wants to merge 6 commits into
InferMatrixLab:mainfrom
ycjcx123:refactor/m7-sampling-module
Draft

ycjcx123 wants to merge 6 commits into
InferMatrixLab:mainfrom
ycjcx123:refactor/m7-sampling-module

Conversation

@ycjcx123

Copy link
Copy Markdown
Contributor

1. What outcome are you trying to achieve?

Verify the correctness of our M7 sampling/RNG refactor on GPU.

2. What is the highest-leverage variable for that outcome?

The correctness of M7 sampling and RNG behavior on real GPUs.

3. How will this work move that variable rather than create busywork?

Run CUDA tests through the interfaces to validate fixed-seed behavior, RNG
progression, preemption/restoration and cleanup, and retain reproducible
results.

4. Who is the customer or user, and what do they actually need?

Maintainers and future developers need reproducible evidence to assess
refactor correctness and validation coverage.

PR necessity and bugfix evidence

A PR is needed because the earlier M7 refactoring work lacked GPU validation.
This test-only follow-up to #80 adds reusable CUDA regression coverage and
records the results.

Production code is unchanged. Failures found in the wider selection are
documented below for separate investigation. The failing test functions,
including their decorators and bodies, are unchanged by this PR. Other edits in
these files only add tests, device-parameterize existing tests, or adjust shared
test helpers; they do not change the failing functions.

Purpose

Validate CUDA sampling distributions, request-local RNG advancement and
cleanup, ordinary asynchronous random sampling, fixed-seed
preemption/restoration, BF16 mixed recurrent execution, and worker/request-ID
lifetimes.

The tests follow the loop/stage architecture and wavefront chunked-prefill
scheduling. Existing async greedy-speculation and CUDA Graph regressions were
also exercised.

Async speculative temperature/top-k/top-p sampling remains a separate follow-up
PR. This PR does not claim full M7 3/3 completion while FA4 and PD remain
unverified.

What changed

  • 8 test files, 6 commits on top of 66920da (fb8b8af..989a20b).
  • Adds or device-parameterizes 43 GPU cases covering the contracts above.
  • No files under vllm_rlt/ change; no docs, defaults or production behavior
    change.

Changed files:

  • tests/test_sampler.py
  • tests/test_engine.py
  • tests/test_speculative.py
  • tests/test_stage_engine.py
  • tests/test_async_pipeline.py
  • tests/test_prefix_growth.py
  • tests/test_split_prefill.py
  • tests/test_mixed_execution.py

Performance and model-quality analysis are not applicable: this PR changes tests
only and does not alter a production execution path.

Test Plan

vllm-rlt Version / Commit:

  • Candidate: 989a20bde9856a9914745b3a0485c635f6057d6d
  • Baseline: 66920da9c1880227707b5d65c6d88bfb4a525ad9
  • Hardware: 1×RTX 5090, 32 GB, SM120
  • Environment: Python 3.12.3, PyTorch 2.8.0+cu128, Triton 3.4.0, pytest 8.4.2
  • Complete pre-commit: CI-pinned Ruff 0.13.2, pre-commit 4.7.0

Source checksums matched before and after testing.

Evidence and checksum

Download the ZIP provided as a PR attachment:

m7-validation-989a20b-public.zip

SHA256:

9fc52a231d64a4642613dbef11c829e7259aa918e12cbc5f74cb6d95e6822ce1

It contains environment/dependency records, JUnit XML, original failure logs,
exact commands and exit codes, baseline diagnostics, the 43-case selection,
unchanged-function comparisons, and complete CPU/pre-commit output.

Reproduction commands and 43-case selection

Run from the candidate repository checkout. The reserved GPU value used for
these runs is recorded in environment.json; export it before GPU commands:

export CUDA_VISIBLE_DEVICES="${M7_GPU:?set M7_GPU to the reserved device from environment.json}"

Verify and extract the evidence ZIP:

echo '9fc52a231d64a4642613dbef11c829e7259aa918e12cbc5f74cb6d95e6822ce1  m7-validation-989a20b-public.zip' \
  | sha256sum --check -

git checkout 989a20bde9856a9914745b3a0485c635f6057d6d
unzip m7-validation-989a20b-public.zip -d m7-evidence

43-case selection: m7-gpu-nodeids.txt lists the exact GPU node IDs
belonging to added/modified top-level test functions. Functions were identified
by comparing complete function ASTs, including decorators, against baseline
66920da; selection was not based on pass/fail outcomes.

m7-case-classification.json records the original JUnit statuses. The included
select_m7_gpu_cases.py can regenerate the list from a checkout containing both
commits.

Run the exact 43 nodes directly through pytest:

python - <<'PY'
from pathlib import Path
import subprocess
import sys

nodes = Path("m7-evidence/m7-gpu-nodeids.txt").read_text().splitlines()
assert len(nodes) == len(set(nodes)) == 43
raise SystemExit(subprocess.call([
    sys.executable, "-m", "pytest", "-q", "--run-gpu",
    "--junitxml=m7-gpu.xml", *nodes,
]))
PY

Reproduce the original wider GPU selections:

M7_TESTS=(
  tests/test_sampler.py
  tests/test_engine.py
  tests/test_speculative.py
  tests/test_stage_engine.py
  tests/test_async_pipeline.py
  tests/test_prefix_growth.py
  tests/test_split_prefill.py
  tests/test_mixed_execution.py
)
FILTER='not fa4 and not flash_attn and not packed_metadata_supports_query_lengths_and_per_request_depths'

python -m pytest "${M7_TESTS[@]}" \
  -m gpu --run-gpu -q --tb=short -ra \
  -k "$FILTER" --junitxml=gpu.xml

python -m pytest \
  tests/test_async_speculative.py \
  tests/test_async_speculative_lifecycle.py \
  tests/test_async_speculative_graphs.py \
  tests/test_input_preparation.py \
  tests/test_cuda_graph.py \
  -m gpu --run-gpu -q --tb=short -ra \
  -k "$FILTER" --junitxml=extra.xml

FILTER is a name-based workaround, not a completeness claim:

  • packed_metadata_supports_query_lengths_and_per_request_depths is an
    FA4-only test whose name does not contain FA4.
  • The two FA4-only metadata cases in test_prefix_growth.py also have no FA4 in
    their names; they escaped this filter and are the F1/F2 failures below.

Reproduce complete CPU regression and pre-commit:

python -m pip install 'ruff==0.13.2' 'pre-commit==4.7.0'

CUDA_VISIBLE_DEVICES='' \
PYTEST_ADDOPTS='-m "not gpu" --tb=short -ra --junitxml=full-cpu-junit.xml' \
pre-commit run --all-files --verbose

Equivalent standalone full CPU selection:

OMP_NUM_THREADS=1 python -m pytest \
  -m 'not gpu' -q --tb=short -ra

Test Result

Selection Result
Added/modified M7 GPU cases 43 passed, 0 failed, 0 skipped
Initial wider GPU selection 117 passed, 4 failed, 280 deselected
Existing async-speculation / Graph / input-preparation selection 73 passed, 1 skipped, 60 deselected
Initial eight-module CPU regression 272 passed, 129 deselected
Complete non-GPU CPU selection 836 passed, 34 skipped, 378 deselected
Complete pre-commit with Ruff 0.13.2 Passed

The 43 M7 cases are a subset of the wider GPU selection. The CUDA sampler smoke
is also a subset and is not additive.

The complete CPU result was produced by the pre-commit pytest hook, not a second
independent run. Its skips concern unavailable optional tokenizers, checkpoints
and lm_eval; exact reasons are retained in the evidence bundle.

Known wider-selection failures and unresolved flake

All three failing test functions, including decorators and bodies, are unchanged
between baseline and candidate. Production files, tests/helpers.py and
tests/conftest.py are also unchanged. See failure-body-comparison.json.

Of the failures rerun on the baseline, only F3 reproduced. F1/F2 were not rerun
because the unchanged capability guard rejects SM120 before kernel execution;
F4 passed baseline diagnostics. Unchanged test bodies alone do not establish
F4's root cause.

Exact failure node IDs and triage

F1

tests/test_prefix_growth.py::test_prefill_uva_metadata_matches_reference_and_waits_for_consumer[False]

F2

tests/test_prefix_growth.py::test_prefill_uva_metadata_matches_reference_and_waits_for_consumer[True]

F3

tests/test_split_prefill.py::test_cuda_split_prefill_survives_decode_graph_capture_and_replay

F4

tests/test_stage_engine.py::test_stage_client_cuda_async_generation_and_close[cuda-False]
Failure Observed Expected Baseline evidence Suggested ownership / next step
F1, F2 Explicit FA4 construction rejects SM120. Metadata validation requires supported paged-FA4 hardware. Exact baseline cases were not rerun; the guard and test bodies are unchanged. The baseline reaches the same guard on SM120. Attention/backend test owner: confirm capability-based gating separately; retain the runtime compatibility guard.
F3 num_preluded_tokens == 4. Existing assertion expects 2. Same failure reproduced on unmodified 66920da, in the eager run before its Graph run. WCPB/split-prefill owner: confirm staging semantics, then revalidate the entire eager/Graph test.
F4 drain() exceeds its 500-iteration limit. Requests finish within the test bound and match local output. Candidate warm rerun passed; baseline warm and isolated cold-cache runs passed. StageClient/test-lifecycle owner: investigate timing or stall evidence. Cause remains unresolved; a passing rerun is not a fix.

F1/F2 fail inside the intentional capability guard in
vllm_rlt/kernels/flash_attention.py::select_version, before any kernel runs.
The unmodified baseline reaches the same guard on SM120.

F4 is not attributed to this PR: it did not reproduce on baseline diagnostics or
a warm candidate rerun. The first failure is retained, and we propose a separate
wall-clock/diagnostic fix or flake issue rather than treating the rerun as a fix.

The extended PD case skipped because NIXL was missing. This single-GPU host also
cannot satisfy its two-GPU requirement.

FA4, PD/NIXL, full-model quality and historical sampling-cost/performance
qualification remain unverified.

Review and additional coverage

@LyxWxj, could you confirm the WCPB staging expectation for F3?

@bjf-frz, could you help triage F4 and confirm the intended capability-gating
policy for F1/F2?

Could anyone help run the applicable subsets on supported paged-FA4 hardware or
a same-host two-GPU NIXL setup? Please advise whether these are required before
marking M7 3/3 ready.

Decisions needed

  1. Is the SM120/Triton evidence (43/43 in-scope GPU cases) sufficient for M7 3/3,
    with FA4 and PD/NIXL listed as unverified?
  2. @LyxWxj: is num_preluded_tokens == 4 the intended WCPB contract for F3?
    If yes, we will fix the stale assertion in a separate PR and rerun the Graph
    branch.
  3. @bjf-frz: is a separate fix for F4 acceptable, and should F1/F2 be
    capability-gated in tests rather than in production?
  4. If FA4 or PD are required before merge, who can provide an SM90/SM100 host
    for FA4 or a same-host two-GPU NIXL setup?

This remains a draft while the remaining qualification requirements are
clarified.

@LyxWxj

LyxWxj commented Oct 26, 2026

Copy link
Copy Markdown
Contributor

Yes, num_preluded_tokens == 4 is intended here. With D=4, prefill_chunk_size=2, and a 4-token batch budget, one WCPB prelude batch can stage both chunks. In general, it can stage up to D chunks per request, bounded by the available token budget, the WCPB window, and the remaining prompt length. max_num_seqs limits the number of distinct requests in the batch. At this point, num_prefilled_tokens remains 0 because neither chunk has completed the final recurrent depth. The assertion expecting 2 is stale for this WCPB test. Could you please update it and rerun the full eager and CUDA Graph capture/replay test in the separate PR? The current failure exits before reaching the Graph portion. Thanks, and apologies for the stale assertion : ).

@ycjcx123

Copy link
Copy Markdown
Contributor Author

Yes, num_preluded_tokens == 4 is intended here. With D=4, prefill_chunk_size=2, and a 4-token batch budget, one WCPB prelude batch can stage both chunks. In general, it can stage up to D chunks per request, bounded by the available token budget, the WCPB window, and the remaining prompt length. max_num_seqs limits the number of distinct requests in the batch. At this point, num_prefilled_tokens remains 0 because neither chunk has completed the final recurrent depth. The assertion expecting 2 is stale for this WCPB test. Could you please update it and rerun the full eager and CUDA Graph capture/replay test in the separate PR? The current failure exits before reaching the Graph portion. Thanks, and apologies for the stale assertion : ).

Thanks, that confirms the assertion is stale. I'll change it to 4 and rerun the full eager + CUDA Graph test in a separate PR. I'll aim to follow up tonight with the PR link and logs, and update the M7 F3 entry to reference that fix.

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