[Draft][Quantization][MoE] NVFP4 native expert kernels (raw-layout decode GEMV + grouped prefill) - #44
Draft
01554 wants to merge 16 commits into
Draft
[Draft][Quantization][MoE] NVFP4 native expert kernels (raw-layout decode GEMV + grouped prefill)#4401554 wants to merge 16 commits into
01554 wants to merge 16 commits into
Conversation
…MV and grouped prefill (kernels + tests only) Extracted from the lab expert-tier branch onto upstream main a97dacb. Device-resident raw ModelOpt NVFP4 bank, Triton decode GEMV (wide-load kernel derived from FreeToken nvfp4_fused_moe.py at af71ba43, Apache-2.0, attributed in file headers) and grouped prefill with route alignment; the loader keeps the checkpoint layout and expands per-expert global scales. No backend selection hook yet; CPU tests load the package by file location so they run without the distributed stack. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…nment skip, no file-location shim Per review: the formal contract is pytest with normal imports (GPU paths skip without CUDA); the package lives under vllm.model_executor, so the tests skip when the distributed stack is unavailable instead of loading the package by file location. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
… selection and device-resident experts - NvFp4MoeBackend.NATIVE, mapped from moe_backend="native" (explicit only, not auto-selected); 16-bit activations like Marlin's a16 path. - NativeNvFp4Experts (FusedMoEExpertsModular): raw ModelOpt bank kept on device, decode GEMV for M <= gemv_rows, grouped prefill above; router weights applied in-kernel, finalize is a no-op; no EP. - ModelOptNvFp4FusedMoE.process_weights_after_loading takes the native branch before the Marlin repack (globals expanded per row, input scales dropped). - loader: lab-runtime dependencies removed; tests use normal imports (full environment), UVA smoke case behind --with-uva. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…adapters as FP32 (review) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…h across layers of the same shape (review) Per-layer persistent prefill scratch would cost ~390 MiB per layer at 4096 tokens x top-10 on the 2560/640 model (~18 GiB over 48 layers); layers run sequentially, so one workspace per (device, bank shape, capacities) is enough. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…end=native selection (review) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…, LoRA, bias and SwiGLU parameters explicitly (review) The kernels compute plain SiLU on BF16 activations; configurations they cannot honour fail at backend selection and construction instead of silently computing something else, and quant-config gemm1 parameters are rejected at weight processing. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…ethod's kernel untouched (review) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
… stream, shape), resolved on the forward stream (review) Stream ordering is the exclusivity argument: layers or models sharing a stream cannot overlap, a different stream gets its own scratch, and graph replay reuses the capture stream's scratch. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…modular kernel workspace (review) No process-global buffers: workspace_shapes declares the flat kernel scratch (max of decode and prefill layouts) and apply() carves the Workspace views from workspace2, which the worker's WorkspaceManager provides per execution lane. The experts keep only bank views, an identity row map and a sticky error flag; kernels warm at weight processing. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…l row count (review) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…ported-reason words (review) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…ip guard (full-environment failures must fail) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…WorkspaceManager path (CPU) Builds the modular kernel with the no-DP/EP prepare/finalize around the native experts and runs decode (gemv) and prefill rows through kernel.apply(), so workspace_shapes() -> WorkspaceManager.get_simultaneous() -> apply() carving is covered by the manager's own buffer rather than a hand-allocated scratch. Outputs must equal the direct apply and the manager must hold at least the declared scratch plus output bytes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
… FusedMoEKernel.apply Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
…(order was swapped) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT Signed-off-by: 01554 <24953377+01554@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Upstream extraction, unit 1 of the team plan (notes/upstream-pr-plan-draft-2026-09-09.md). Draft: not ready, not to be merged before two reviews (astra_mid: code/tests; astra: independent diff/deps/GPU).
Base: upstream/main a97dacb (fork branch
lab/upstream-base-20260909). Head: 848f0ef (a452302 plus one tests-only commit: the native experts run through the realFusedMoEKernelwith the no-DP/EP prepare/finalize andWorkspaceManageron CPU). Dependencies: none.Scope: new package
vllm/model_executor/layers/quantization/nvfp4_native/—kernels.py(Triton decode GEMV; wide-load kernel derived from FreeTokenkernel/triton/nvfp4_fused_moe.pyat af71ba43, Apache-2.0, attributed in file headers),bank.py(raw ModelOpt bank validation, workspace,gemvadapter, CPU reference),prefill_kernels.py+prefill.py(grouped prefill with route alignment),loader.py(keeps checkpoint layout, expands per-expert global scales, drops input scales),experts.py(NativeNvFp4Experts, aFusedMoEExpertsModular: decode GEMV for M <= gemv_rows, grouped prefill above, no EP). Selection:NvFp4MoeBackend.NATIVEviamoe_backend="native"only (never auto-selected), a16 quant config like Marlin,ModelOptNvFp4FusedMoE.process_weights_after_loadingnative branch before the Marlin repack;MoEBackendliteral gains"native". Device-resident only. Not in this PR: host/UVA sources, expert cache, runner changes.Executed (Mac, CPU): ruff, ruff format, mypy, pre-commit on all touched files. Not executed locally: the tests (
tests/quantization/nvfp4_native/: kernel oracle tests, loader tests, experts tests) importvllm.model_executornormally and need the full environment; not executed anywhere yet: CUDA runs on this base (kernel smoke, model load with--moe-backend native, graph capture/replay) — pending astra on a clean a97dacb build.Provenance: FreeToken (Apache-2.0) only; no llama.cpp/ggml code in this PR.
Exact upstream/main a97 wheel (astra, 2026-09-09): a452: CPU 36 passed + 4 subtests, decode 4, grouped 2 + decode 4 PASS; 848f0ef: CPU 37 passed + 4 subtests, 0 skipped (results/upstream-native-exact-848f0efca374-20260909). Real-checkpoint smoke on the integration head
lab/upstream-native-ple-mmap0f1a720 (this PR plus the sync PLE base, 96 GiB, no balloon): load, full capture, two requests, same outputs as the Marlin backend on the same head (results/upstream-marlin-checkpoint-0f1a720f3764-20260909). Not a performance comparison; the 48 GiB configuration is not yet run.Kernel benchmark on the exact base (astra, 2026-09-09; benchmarks/nvfp4_native on
lab/upstream-native-bench5e7deba; layer 0, 512 experts, top-k 10, synthetic inputs, CUDA-graph replay 30x20, both route patterns; results/upstream-native-bench-exact-5e7debaa3c00-full-20260909): the Marlin backend is faster than this native backend at every shape where both pass. M=1 decode: Marlin 0.0287 ms vs native (through the real kernel + WorkspaceManager) 0.0476 ms; grouped prefill M=2..1024: native 5-15x slower. This PR therefore makes no speed claim. What it offers is the raw-checkpoint layout with no repack (weights, block scales and per-row globals stay expert-major and host-source friendly for the expert cache), an independent oracle-checked path, and full correctness at all 16 shapes; Marlin had 1-5 elements over the prior tolerance at 3 of 16 shapes (M=256/1024); the differences are consistent with BF16 rounding-order effects (one ulp of intermediate terms, on outputs near zero), the cause is not separated, and this is not evidence of a bug or of quality inferiority.🤖 Generated with Claude Code
https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT