[New Model][Nvidia] Add SM12x support for DeepSeek V4 Flash with essential fixes - #41834
[New Model][Nvidia] Add SM12x support for DeepSeek V4 Flash with essential fixes#41834jasl wants to merge 207 commits into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@zyongye |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements support for DeepSeek V4 on SM12x (Blackwell) architectures by providing Triton-based fallbacks for DeepGEMM-dependent operations. Key enhancements include the introduction of specialized Triton kernels for sparse MLA, FP8 einsum, and MQA logits, as well as memory optimizations in the sparse attention indexer to compute top-k indices without materializing full logits. Additionally, the PR updates the model loader to support weight name filtering for skipping MTP weights and handles Blackwell-specific FP8 quantization scales. I have no feedback to provide.
💡 Codex Reviewvllm/vllm/model_executor/layers/sparse_attn_indexer.py Lines 86 to 89 in 9596dbf  Keep DeepGEMM requirement for SM120 FP4 indexer path
This helper now disables the DeepGEMM requirement for every SM120 run, but the FP4 indexer cache path still depends on DeepGEMM kernels ( vllm/vllm/model_executor/model_loader/default_loader.py Lines 236 to 240 in 9596dbf  Propagate weight_name_filter to fast safetensor loaders
The new pre-load ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
jasl
changed the title
042e366 to
df2e6f8
Compare
test_indexer_builder_keeps_short_prefill_continuations_as_prefills builds the metadata builder via object.__new__ and set reorder_batch_threshold, but build() reads self.decode_threshold (renamed upstream) -> AttributeError. Pre-existing red before this merge; the upstream PCP merge (vllm-project#46570) additionally makes build() read self.use_pcp. Set both on the mock. use_pcp=False keeps the assertion validating the DSv4 ubatch guard (treat_short_extends_as_decodes == not has_prefilling_rows).
Add transport-neutral compact accounting and address types, a fixed-page scatter allocator, and variable-size eviction support for the built-in LRU and ARC policies. The legacy CachePolicy.evict extension point remains unchanged. Compact callers use a new fail-loud evict_until contract, while unsupported external policies remain instantiable until compact mode is requested. This commit contains no manager, worker, transfer, shared-memory, rank- ownership, bounded-tail, or model-specific runtime behavior. Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Carry an optional two-field compact signature through KVCacheConfig and the normalized offloading configuration. The signature records group identity and real payload bytes per native GPU block per worker; blocks_per_chunk remains a separate runtime dimension. Worker-side packed tensor metadata is converted into immutable per-slice geometry and checked against the scheduler-safe signature. Compact behavior is disabled when enable_compact_layout is absent, and pipeline parallelism greater than one fails explicitly. Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Add a pure direction-neutral planner that maps actual GPU block IDs and packed slice offsets to compact CPU addresses, including fragmented physical spans and partial first chunks. The planner validates group ordering, bounds, slice payload units, full GPU ID and CPU address consumption, and exact per-address payload geometry. It returns immutable pointer and size descriptors without performing DMA. Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Extend the existing CPUOffloadingManager with an optional compact path backed by one global FixedPageAllocator coordinate space. Store preparation performs atomic batch-fit simulation and variable-size eviction, then returns complete fragmented physical spans before DMA. Pending and committed allocations retain the canonical lifecycle. Loads own BlockStatus references and are non-evictable while active; completion restores evictability. Write-pending and evictable counters remain exact across store success, failure, load, eviction, and reset. The legacy manager path is unchanged when compact geometry is absent. Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Activate the compact path through the existing normalized config, CPU spec, manager, worker, and asynchronous descriptor submission lifecycle. The spec uses fixed 64 KiB scatter pages reduced by the per-rank budget GCD. The worker allocates one rank-local pinned region and executes planner descriptors for both GPU-to-CPU stores and CPU-to-GPU loads. Compact payload units are per native GPU block per worker and are multiplied by blocks_per_chunk exactly once. Tests exercise the real compact transfer method in both directions with packed slice offsets, actual GPU block IDs, fragmented CPU spans, partial chunks, pointer reversal, and exact byte counts. The legacy path remains inert when compact geometry is absent. Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
Signed-off-by: Colton Ottley <colton@ottleyengineering.com>
|
Submitted as jasl#32 against The six-commit series passed the focused source suite plus exact-image hardware qualification on 2x RTX PRO 6000 Blackwell GPUs at TP=2 and 4x RTX PRO 6000 Blackwell GPUs at TP=4. The PR body includes the fixed-page compact-layout results, the hardware-only integration defects found during qualification, native CUDA copy tests, and the final semantic CPU-hop evidence. |
External contribution by coltonottley. Optional compact CPU KV offload layout for packed heterogeneous caches; disabled by default (kv_connector_extra_config. enable_compact_layout=true). Validated on 2x RTX PRO 6000 Blackwell TP=2, 223 tests. Reviewed (4-lens adversarial): opt-in isolation holds — the default path (compact off, PCP=1) is byte-identical; no interaction with our vllm-project#48959/vllm-project#48911 offload deltas. One compact-ON init bug (spec.py gcd page-size collapse) fixed in the following commit.
spec.py derived the compact page size from gcd(64 KiB, per_rank_budget). A budget with low 2-adic valuation (e.g. a round decimal cpu_bytes_to_use like 100 GB, whose per-rank value factors as 2^10 * odd) collapses the page size to ~1 KiB or less, so FixedPageAllocator eagerly builds a free-page list of tens of millions to billions of entries and OOM/hangs at worker init. Reachable only with enable_compact_layout=true. Round the per-rank budget down to a whole number of fixed 64 KiB pages instead (its own design intent — 'stable 64 KiB fixed-page target'), so the page size always divides the budget, the free-page count stays bounded (budget/64 KiB), and < 64 KiB/rank is wasted. Reject a sub-page budget with a clear error rather than constructing a zero-page allocator. Found in review of PR #32; inert for GB10 (never enables CPU offload). Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…date page test PR #32 multiplied build_offloading_config's tokens_per_block by prefill_context_parallel_size (context_parallel_factor = dcp * pcp). That breaks the upstream invariant that KV-cache sharding ignores PCP (parallel.py documents PCP does not increase shard count), which the upstream test test_offloading_spec_kv_sharding_ignores_prefill_context_parallel (added by vllm-project#46570) asserts (tokens_per_block == (16,) at PCP=2, not (32,)). Restore the upstream computation (block_size * decode_context_parallel_size). No-op on GB10 (PCP=1). Also update test_spec_compact_page_size_non_divisible_budget to assert the new budget-rounding page-size behavior instead of the removed gcd behavior. Both found by running PR #32's full suite (not just the compact tests) post-merge. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
|
Heads-up for anyone benchmarking DP topologies on this branch: DeepSeek-V4-Flash (stock MXFP4) with The root cause is on main (#42118), not in this PR; fix in #49335, validated fail-to-pass on 8x RTX 5090, 4x RTX PRO 6000 and 2x B200, including on this branch. Until it lands here, MARLIN is a safe MoE backend under DP, or apply the small diff from the PR. |
When the FlashInfer kernel is stable, I shall re-evaluate the latest upstream changes and migrate to the faster path |
…20260726 # Conflicts: # requirements/cuda.txt # vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py # vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py # vllm/envs.py # vllm/model_executor/layers/fused_moe/fused_moe.py # vllm/models/deepseek_v4/common/ops/cache_utils.py # vllm/models/deepseek_v4/nvidia/dspark.py # vllm/models/deepseek_v4/nvidia/ops/o_proj.py # vllm/v1/attention/backends/mla/sparse_swa.py # vllm/v1/core/kv_cache_manager.py # vllm/v1/kv_offload/config.py # vllm/v1/worker/gpu/model_runner.py # vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
_fp8_paged_mqa_logits_rowwise_kernel is the DSv4 sparse-indexer decode logits
kernel on SM12x. The decode warmup dummies run with seq_lens == max_query_len
(3 for MTP2), so DeepseekV4Indexer.forward short-circuits on
max_seq_len // compress_ratio <= topk_tokens and never reaches the paged-MQA
path; the kernel therefore JIT-compiles inside the first request whose context
exceeds compress_ratio * index_topk, costing a latency spike mid-inference
(now visible via upstream's new jit_monitor).
Warm the 6 reachable compile variants (num_rows in {==1, non-16-aligned,
16-aligned} x logits_width in {16-aligned, not}). Fidelity details that matter:
the KV strides are read off the BOUND cache tensor (a padded strided view with
page stride 8640, not a contiguous 64*132=8448) and the block-table row width
from the live indexer metadata builder, because both are constexprs -- a
synthetic contiguous cache would compile a different cubin and leave the real
one to JIT anyway. next_n is 1 on SM12x even under MTP2 because the indexer
builder flattens next_n not in (1, 2).
Guarded by try/except: warmup must never break startup.
Not covered here (deliberate): _tf32_hc_prenorm_gemm_kernel has ~68 variants
(SM-count-derived NUM_SPLIT), too expensive to warm on every boot for a
cold-cache-only benefit; _compute_global_topk_indices_and_lens_kernel needs
pointer-alignment-class construction that rests on assumed Triton internals.
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
- Make quant_config an instance attribute for reuse across components - Add pad_shared_expert flag enabled when weight_block_size is set and sequence_parallel_moe is disabled - Remap .w2 -> .down_proj in shared_experts during weight loading - Pad shared expert weights via DeepseekV4Model._pad_shared_expert_weight to align with block quantization constraints Signed-off-by: Hunter L. Allen <hunter@allen-software.com>
Merges #27 (+1523/-466, 23 files) onto 520c259. The PR was based on d15fad8 (302 commits back) yet produced only ONE conflict. Conflict resolution (rejection_sampler.py __call__): took OUR chunked structure -- upstream since refactored _verify_in_chunks to own logprobs internally, so the PR's old-signature call site is obsolete -- while threading through the PR's actual contribution, draft_logits_index_mapping. ★ Also fixes a SILENT MIS-BINDING the automerge introduced (not in any conflict): rejection_sample() gained draft_logits_index_mapping as positional parameter #12, immediately BEFORE synthetic_conditional_rates. _verify() passed self.synthetic_conditional_rates positionally at slot #12, so after the merge it bound to draft_logits_index_mapping and synthetic_conditional_rates silently became None -- corrupting spec-decode sampling with no error. Both trailing optionals are now passed by keyword, and the hazard is called out in a comment. Threading verified end-to-end: producer model_runner.py:1172 -> __call__ -> _verify_in_chunks -> _verify -> rejection_sample (consumer rejection_sampler_utils.py). Without it the parameter would be accepted and silently dropped, disabling the feature. Swept every other rejection_sample call site: v1/sample/ and kernel_warmup use a different function and pass all-keyword; no other mis-binding. Co-authored-by: alexbi29 Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
|
This pull request has merge conflicts that must be resolved before it can be |
# Conflicts: # tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
…s_and_lens The top-k indices fed to this kernel are data, not loop bounds: they come from the indexer's top-k, which writes into a torch.empty buffer shared by every layer. An unwritten or corrupted slot can hold a large positive value, and the kernel validated only `local_idx >= 0`, so such a value gathered past the end of the block table -- an illegal access that takes down every TP rank. Bound the row offset as well. Out-of-range entries fall out as -1 and are excluded from topk_lens, so a producer bug degrades one token instead of killing the engine. The bound is structural (block_indices < block_table_stride keeps the gather inside the request's own row) and cannot reject a legitimate index: the largest valid compressed position maps to num_blocks-1, and num_blocks is at most max_num_blocks_per_req, which is the stride. Not reachable on our hardware today -- SM12x asserts use_fp4_indexer_cache off, so q_scale is always None, so prefill always takes the fused DeepGEMM top-k path that never materializes logits. Upstream hit it on stock vLLM when its SM12x logits kernel began emitting NaN bit patterns as indices (vllm-project#49896, root cause 3). Kept as local hardening rather than an upstream PR, and deliberately minimal: one term plus a contract test, easy to drop when the DSv4 sparse path is reimplemented against the newer upstream/FlashInfer primitives.
|
This pull request has merge conflicts that must be resolved before it can be |
|
Independent 2× GB10 run (SM121, TP=2, 200GbE RoCE) of the validated tag
Lines up with your reported numbers for this tag (ctx_pp 1757/1774/1699). Vs our prior early-July build ( Confirming your CUDA-graph-profiler note: at 256K long-context still solid: a 198K-token needle returns exact retrieval end-to-end (via our LiteLLM gateway → served model). Serve profile unchanged from our earlier posts (MTP2, mml 262144). Thanks for cutting the stable tag — made for a low-risk build target. Update — we rolled this build back to our prior head ( |
20 participants
Summary
This PR enables DeepSeek V4 Flash on SM120/SM121 Blackwell client hardware by carrying the SM12x fallback and tuning stack needed for the current vLLM V1 path. It is intended for RTX PRO 6000 Blackwell Workstation Edition, RTX 5090-class SM120, and GB10 / DGX Spark SM121 users who cannot use SM100-only TMEM /
tcgen05kernels.This branch is reconciled on top of the merged #43477 and provides the stock-deps path: DeepSeek V4 on SM120/121 that builds and serves on released FlashInfer / DeepGEMM wheels, complementing #43477's route that needs the unreleased FlashInfer #3395 + DeepGEMM #324 dependency branches. It is kept synced onto current
upstream/mainand now also carries DSpark spec-decode support (self-drafting block-5) alongside the default MTP2 path. Latest validated head is tagsm120-pr-41834-stable-preview-20260711(b5c0d43b96) — see Update 2026-07-11 below.Change footprint — model kernels vs. core-vLLM touch points
The branch splits cleanly into model/kernel code and a small set of core-vLLM integration points, of which ~+3.5k is tests:
vllm/models/deepseek_v4/**plus the SM12x sparse-MLA decode / indexer / DeepGEMM kernels that live in shared dirs (v1/attention/backends/mla/sparse_mla_kernels.py,model_executor/layers/sparse_attn_indexer.py,v1/attention/backends/mla/{indexer,sparse_swa}.py,utils/deep_gemm.py,kernels/mhc/tilelang.py), the new DSv4 reasoning parser / tokenizer, and device tuning JSONs.models/deepseek_v4/sparse_mla.py, perf) —_c128a_effective_topk_widthtakes the max position from the CPU-sideCommonAttentionMetadata.max_seq_leninstead of a per-stepint(positions.max().item())device sync, dropping a launch-stream stall on every C128A metadata step. Decode is identical (max_seq_len-1 == positions.max()); only chunked prefill sees a safe, slightly-wider 128-aligned top-k.single_type_kv_cache_manager.py,kv_cache_coordinator.py,kv_cache_manager.py,sched/scheduler.py(+1)cache_blockstail-block-reuse rewritev1/spec_decode/llm_base_proposer.pyfused_moe.py,oracle/mxfp4.py,routed_experts.py,experts/flashinfer_cutlass_moe.py,quantization/mxfp4.py,oracle/nvfp4.pyquantization/utils/fp8_utils.py,linear/scaled_mm/{cutlass,marlin}.py,csrc/.../marlin_moe_wna16/ops.cu(the only C++)config/vllm.py,compilation/breakable_cudagraph.py,passes/utility/fix_functionalization.py,config/compilation.pychat_completion/protocol.py,serve/render/serving.py,tool_parsers/structural_tag_registry.py,chat_utils.py,engine/protocol.py,chat_completion/{serving,batch_serving}.py,reasoning/__init__.pyreasoning_content/thinkingparam / tool-call streaming (jasl#19 instruction-following)model_executor/warmup/kernel_warmup.pyweight_utils.py,default_loader.pyenvs.py,utils/flashinfer.py,utils/import_utils.py,v1/worker/{gpu_model_runner,ubatch_utils}.pyVLLM_DEEPSEEK_V4_*flags +has_cutedsl/has_flashinfer_trtllm_sparse_mlaprobesTwo notes for review:
kv_cache_coordinatorcache_blocksrewrite (affects hybrid-KV models; validated ≥ prior behavior), the MTP proposer base-class change, and the OpenAI-entrypoint plumbing. Everything else (MoE oracle, fp8_utils, cudagraph gate, warmup, envs) is arch / quant / env-gated and inert for other models.Duplicate-work check
The nearest open/merged PRs are related but not duplicates:
42657aca65) and carries the stock-deps DSv4 SM120/121 path that runs on released wheels, complementing #43477's fork-deps route.compress_ratio=0unscaled-draft-rope fix — carried in this branch as of 2026-07-11 (see below).Fixed preview tags
These tags are in
jasl/vllmand give users stable pins while the PR is still moving:sm120-pr-41834-stable-preview-20260711b5c0d43b96upstream/main(181 commits), #48304 MTP unscaled-draft-rope, ~1097-line dead-kernel cleanup. GB10: GSM8K 0.96, tool-call 42/42, arthur c=1 4/4; deep-prefill +6-10%. See Update 2026-07-11.sm120-pr-41834-stable-preview-202607059f18be7630sm120-pr-41834-stable-preview-20260704b43470e871constexpr→runtime (stops the Triton recompile → unified-memory leak → hard-freeze) + fp8-einsumtl.multiple_of(16)(~24% decode @256k).sm120-pr-41834-stable-preview-20260703444fe3ac8bpersistent_topkfor <128 KB-smem parts (GB10 ≥400k long-context).sm120-pr-41834-stable-preview-20260626c766cbc6ffupstream/main+ C128A metadata device-sync removal; our NVFP4 CUTLASS-clamp fix landed upstream as #46492.Older tags (
…-20260623back to…-20260612…) remain injasl/vllmfor history.Update 2026-07-11 — synced onto upstream/main (181 commits) + DSpark MTP-rope fix; fresh baseline
Latest validated head
b5c0d43b96(tagsm120-pr-41834-stable-preview-20260711), 181 upstream commits past the 07-05 tag.What's in it
upstream/main(merge of3d99b0499a; 7 conflicts resolved, no fork feature dropped). Inherits the scheduler max-length #447 KV-cache-spec refactor, the libtorch-stable custom-ops ABI migration (vllm._C→vllm._C_stable_libtorch), and [Perf] Cachetoken_to_req_indicesfor dsv4, 5x~6x kernel performance improvement #47474 dsv4token_to_req_indicescaching.routed_scaling_factornow applied inside the topk kernel ([Kernel] Applies routed_scaling_factor internally #47408) — verified single-application, no double-scale.compress_ratio=0entry (unscaled draft rope) — port of [Bugfix] Honor DeepSeek V4 checkpoints' MTP-layer compress_ratio entry (unscaled draft rope) #48304: DSpark-style checkpoints whosecompress_ratiosinclude the MTP draft layer with entry 0 get uncompressed-KV + plain (non-YaRN) rope for the draft instead of forced yarn-scaled rope. ([Bugfix] Complete FlashInfer CUTLASS MXFP4 MoE wiring for DeepSeek-family models #48303's FlashInfer-CUTLASS MXFP4 MoE wiring was already carried here.)Validation (GB10 SM121, 2-node TP=2; DeepSeek-V4-Flash + MTP2, FlashInfer 0.6.14, nccl 2.30.7)
Correctness:
tool_choice+ thinking, 42 reqs): 42/42, 0×500Perf — pinned llama-benchy standard (fp8 KV, prefix-cache on,
FULL_AND_PIECEWISE, mml 49152, util 0.85; pp2048 tg128, C=1, 3 runs):Deep-context prefill up vs the 07-05 baseline (d32768 pp2048 +10%, ctx_pp +6.5%); decode healthy at all depths.
Dependency-behavior note: upstream's CUDA-graph memory profiler (default since v0.21.0) now reserves cudagraph + peak-activation memory before sizing KV, so KV cache at util 0.85 is ~171k tokens (vs ~284k on the older build, which over-allocated). Restore the prior effective KV size with
--gpu-memory-utilization 0.876orVLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0. No effect on the C=1 numbers above — it only bounds high-concurrency long-context capacity.Running DSpark
DSpark is DeepSeek's self-drafting speculative-decode variant (draft weights carried in the target checkpoint, block size 5). Serve the DSpark checkpoint with the
dsparkspeculative method:vllm serve deepseek-ai/DeepSeek-V4-Flash-DSpark \ --tokenizer-mode deepseek_v4 --trust-remote-code \ --tensor-parallel-size 2 --kv-cache-dtype fp8 \ --speculative_config '{"method":"dspark","num_speculative_tokens":5}'num_speculative_tokensmust be 5 (the checkpoint'sdspark_block_size); the draft is self-hosted, so no separate--speculative-modelis needed.VLLM_USE_V2_MODEL_RUNNER=1; V2's long-context recall is correct after the Add CUDA graph-based all reduce launcher #26 padded-Q fix.--gpu-memory-utilization/--max-model-lenaccordingly.Dependencies (stock-deps path)
requirements/cuda.txt(flashinfer-python+ the GitHub-releaseflashinfer-cubin); it ships the SM120 packed sparse-MLA kernels, so a stock build picks it up — no manual FlashInfer install dance.nvidia-nccl-cu13==2.30.7per node (a rebuild reverts it to torch's 2.28.9; a per-node mismatch hangs the NCCL handshake).VLLM_DEEPSEEK_V4_FLASHINFER_SM120_DECODE) and prefill (VLLM_DEEPSEEK_V4_FLASHINFER_SM120_PREFILL) FlashInfer sparse-MLA paths now default on; set either=0to fall back to the FlashMLA / Triton path (a no-edit escape hatch, availability-gated so stock installs without the kernel degrade gracefully rather than raising).Running the NVFP4 checkpoint
This branch also serves
nvidia/DeepSeek-V4-Flash-NVFP4on SM12x (RTX PRO 6000 / GB10). The NVFP4 MoE auto-selects the FlashInfer CUTLASS backend (the SwiGLU-clamp model gate accepts it), so no--moe-backendflag is required, and no special FlashInfer build is needed:--kv-cache-dtype fp8is mandatory: DeepSeek-V4'sfp8_ds_mlaattention asserts an fp8 KV layout, so the defaultautofails at model construction (this is not NVFP4-specific). Expert-parallel off (plain TP) is the supported path.Accuracy matches MXFP4 (GSM8K 8-shot ~0.96 on both SM120 and SM121). On SM12x NVFP4 is not a memory or throughput win versus MXFP4: NVFP4 weights are ~4 GiB/GPU larger (~78 vs ~74 GiB), leaving less KV-cache room; single-stream prefill is marginally faster and aggregate decode marginally slower. Its value here is checkpoint availability / parity with the SM100 datacenter path — MXFP4 remains the better practical choice on consumer Blackwell.
AI assistance disclosure
AI assistants, including OpenAI Codex/GPT models and Anthropic Claude models, were used for code review, refactoring support, regression-script writing, and benchmark analysis. The branch was validated through human review plus the commands and harness artifacts listed above.