[JJ] Preallocate DeepSeek V4 padded-query output

[voipmonitor]

Status

Implemented and TP2-qualified for DeepSeek-V4-Flash-0731 with fixed probabilistic DSpark K5 on NVIDIA SM120.

Operation contract

The fused DeepSeek V4 query-normalization, RoPE, and KV-insert operator writes a padded query tensor consumed by sparse MLA. The model runner must own this output and include its peak size in memory admission before vLLM assigns remaining device memory to GPU KV blocks.

Resulting behavior

• The stable-ABI operator accepts a caller-owned q_out tensor and performs no hidden output allocation.
• V1 and V2 model runners reserve reusable padded-query storage during memory profiling.
• V1 DBO profiles every scheduler-reachable microbatch workspace.
• NVIDIA and AMD DSpark context-KV insertion use the same mutating operator signature.
• Functionalization derives mutated arguments from the registered schema and supports both the caller-owned output contract and the legacy in-place query contract.
• Operator validation rejects incompatible shapes, dtypes, layouts, and devices before kernel launch.

The operator API changes with this PR. The Python runtime and compiled vLLM extension must be updated together; out-of-tree callers must provide q_out.

The caller-owned output design originates from jasl/vllm#26. This PR adds V1/DBO reservation, functionalization coverage, and both DSpark context-KV call sites required by the Jovian Judgement runtime.

Validation

Registry artifact:

voipmonitor/vllm:jovian-judgement-vllma67b59a-b12xaa76f04-fi803c466-cu133-torch213-20260905-r6
sha256:8222ac5d319c0f4dae04e6a2abd379745c6a2433b1c5c7b454c9fed076d84b08

Source and runtime conditions:

• vLLM base: dev/jovian-judgement@b7e3d033676d5db46fb7d6cdd40d760365a1e239
• vLLM integration tree: a67b59a4099457fbcdadce4476c88504fafaf083
• CUDA 13.3, PyTorch 2.13, RTX PRO 6000 Blackwell, TP2/DCP1
• DeepSeek-V4-Flash-0731, fixed probabilistic DSpark K5, MNS8, MNB4096
• FULL_AND_PIECEWISE CUDA graphs

Results:

• DSpark NVIDIA/AMD context-KV helper tests: 2 passed.
• Focused Ruff check and format: passed.
• Target capture: 9/9 PIECEWISE and 7/7 FULL shapes.
• DSpark capture: 6/6 FULL shapes.
• GPU KV admission completed with 1,199,317 tokens after padded-query workspace accounting.
• A 144,028-token engine-driven LMCache cold/restore run restored 143,360 external tokens and matched every retained byte across eight heterogeneous KV groups on both TP ranks.
• Sustained CC1: 201.1 output tok/s and 72.6 engine steps/s.
• No runtime allocation, graph-capture, or service error occurred during qualification.

Machine-readable receipt: https://github.com/local-inference-lab/blackwell-llm-docker/blob/main/validation/jovian-judgement-ds4-r6-engine-driven-lmcache.json

Compatibility

No model arithmetic, attention selection, KV layout, scheduler policy, or sampling behavior changes. The reserved workspace reduces memory otherwise available for GPU KV blocks and prevents an unaccounted allocation from first appearing under live traffic.

Development disclosure

The implementation and validation were completed with OpenAI Codex assistance under human direction.

[coderabbitai]

Warning

Review limit reached

Next included review available in 59 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: 6b998c71-227e-4c3c-82e1-bb6a0e834b44

📥 Commits

Reviewing files that changed from the base of the PR and between b7e3d03 and 6efc116.

📒 Files selected for processing (12)

• csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu
• csrc/libtorch_stable/ops.h
• csrc/libtorch_stable/torch_bindings.cpp
• tests/compile/passes/test_functionalization.py
• tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
• tests/models/test_deepseek_v4_dspark_context_kv.py
• vllm/compilation/passes/utility/fix_functionalization.py
• vllm/models/deepseek_v4/amd/dspark.py
• vllm/models/deepseek_v4/attention.py
• vllm/models/deepseek_v4/nvidia/dspark.py
• vllm/v1/worker/gpu/model_runner.py
• vllm/v1/worker/gpu_model_runner.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[voipmonitor]

Included in dev/jovian-judgement through this PR's individual merge. The reviewed head and its contributor commits remain ancestors; the PR is merged and closed.

Source validation: replaying all 32 R35 review heads on the pinned base exactly reproduces the released Docker's vLLM tree; all 6,870 installed tracked files match. JJ additionally preserves Luke's DS4.1 work and #734. The final composition passed 247 focused checkpoint/scheduler, sampler/warmup and native GPU tests. This is combined-source evidence, not a fresh performance or full-model qualification for this individual PR.

Publication-history clarification: the individual merge linked above is in JJ's first-parent history. It replaces the receipt's archived wrapper-merge reference; GitHub's historical merge SHA may still identify that archive. See #731 for component review order and qualification limits.