[DeepSeek-V4][Preview][DO NOT MERGE] Batch-invariant kernels for 0-diff RL

[ISEEKYAN]

Important

Preview only — DO NOT MERGE.

This draft collects the vLLM changes required for end-to-end DeepSeek-V4
train/rollout alignment. It is intentionally larger than an upstreamable PR.
After the preview is validated, it will be rebased onto current main and
split by kernel, sparse attention, MoE/DeepEP, and reload support. Those
focused PRs are the changes intended to merge. The verl production recipe
will be updated after they land on official vLLM branches.

Purpose

Make DeepSeek-V4 rollout execution independent of request batching so the same
vLLM kernels can be used by the mLite actor recomputation path and produce
exactly matching probabilities.

Preview source:

• head: ISEEKYAN/vllm:ds4-v9-rc1
• revision: 1a6542cdebbaabbf9972c358de8d6baa54e46b0c
• tree: df12131335d84e9848e188905449230902eaf7ca
• fork point used by the image: 7aa248fcfef5ba7a6bfb0ce314e328ce63abb9f9

Default vLLM behavior is unchanged. The aligned path is enabled only with
VLLM_BATCH_INVARIANT=1 and the required native libraries.

Fused activation and FP8 quantization

csrc/batch_invariant/ adds a standalone
fused_silu_mul_per_token_group_quant operator for contiguous and
expert-masked inputs.

• Fuses SiLU, multiply, and 128-element group quantization with fixed launch
  geometry rather than selecting a path from batch-local M.
• Supports FP32, ceil-UE8M0, and packed UE8M0 scales.
• Uses architecture-specific activation/rounding behavior for Hopper and
  Blackwell.
• Uses 64-bit group and output offsets, including tensors above 2**31
  elements.
• Is wired into deep_gemm_moe.py and batched_deep_gemm_moe.py.

In batch-invariant mode, missing kernel support is a hard error rather than a
silent fallback to a numerically different implementation.

Deterministic sparse top-k

csrc/ds4_bi_topk_extension.cu adds ds4_bi.top_k_per_row_prefill.

• Selects top-k from a composite score/index key with a stable tie-break.
• Uses the same selector for prefill and batch-invariant decode.
• Canonicalizes selected indices before FlashMLA so reduction order is fixed.
• Keeps persistent row-start buffers for CUDA graph replay.

This removes top-k set/order changes caused by equal scores, request offsets,
or decode-versus-prefill dispatch.

Sparse MLA and FlashMLA

vllm/models/deepseek_v4/nvidia/flashmla.py and the sparse-attention helpers
add VLLM_DS4_DECODE_KERNEL=sparse.

• Decode gathers compressed and sliding-window KV and calls
  flash_mla_sparse_fwd, matching the prefill attention path.
• Batch-invariant prefill chunks each request independently instead of using a
  grouped batch-local plan.
• C128 top-k width, gather capacity, and workspaces derive from model capacity,
  not neighboring sequence lengths.
• Graph-safe Triton helpers fill top-k and length buffers without host sync.

This makes sparse index selection, KV gather order, and attention reduction
independent of the other requests in the batch.

DeepGEMM, MoE, and DeepEP

• vllm/utils/deep_gemm.py probes and enables the batch-invariant DeepGEMM
  contract, failing closed when the installed fork does not support it.
• MoE backend selection requires a batch-invariant implementation and uses a
  fixed expert capacity instead of expert_num_tokens.max().
• Router gate logits use a common FP32 torch.mm path rather than an
  M-dependent low-latency GEMM.
• DeepEP low-latency dispatch is staged at fixed capacity; empty ranks still
  join collectives and FP8 dispatch requests aligned quantization.
• Shared experts launch once across staged dispatch slices.

Together these changes fix route ordering, expert input quantization, grouped
GEMM geometry, and combine order across actor and rollout.

MHC, CUDA graphs, and weight reload

• Batch-invariant DeepSeek-V4 uses explicit MHC pre/post operations instead of
  the batch-dependent fused path.
• MHC split count is fixed and multi-token inputs are processed in stable
  token order.
• Sparse-attention, top-k, and MHC buffers keep stable addresses across CUDA
  graph capture and replay.
• Layerwise actor-to-rollout reload calls
  finalize_weights_after_layerwise_reload() to refresh MHC broadcast weights.
• attn_sink loading preserves -inf padding through the model weight loader.

Hopper and Blackwell

• Hopper keeps FP32 block scales and uses the compatible block-scaled
  projection path in batch-invariant mode.
• Blackwell keeps packed UE8M0 scales and the DeepGEMM einsum path.
• Both architectures use their validated SiLU and FP8 packing behavior rather
  than treating SM90 and SM100+ as numerically interchangeable.

Test Plan

• Run the batch-invariant kernel contract and large-tensor offset tests.
• Run deterministic sparse top-k tie/order tests.
• Run sparse FlashMLA and DeepSeek-V4 sparse decode tests.
• Run batch-invariant DeepGEMM/MoE and router gate tests.
• Run MHC graph/refit and layerwise reload tests.
• Load both custom libraries on x86_64 and arm64 images.
• Run the companion verl quick_alignment_test, H100 aligned profile, and
  GB200 aligned profile.

Test Result

• image digest: TODO
• verl recipe PR: TODO_VERL_PR_URL
• quick alignment test: TODO
• H100 64-GPU aligned run: TODO
• GB200 32-GPU aligned run: TODO

---

Essential PR description checklist

• Purpose and non-merge preview status are explicit.
• Test plan is listed.
• Test results are attached.
• Documentation and split PR ownership are finalized.

Acknowledgements

The train/inference alignment strategy and deterministic kernel work are
informed by THUDM/slime#2262
and its SGLang/DeepGEMM/DeepEP alignment work led by
@zhuzilin.

Contributors

• Yan Bai (@ISEEKYAN) maintains the DS4 preview
  integration and the batch-invariant DeepSeek-V4 execution path.
• @aoshen02 authored
  36687196d62e0e26d4c29f19d505b1fc4e1df1e8, which widened fused-SiLU
  group counts and output offsets to 64-bit. The fix came from three failed
  eight-node RL runs with long responses, added a numel >= 2**31 regression,
  and verified chunked/full bitwise parity on GB200 without measurable
  regression.
• @andakai and
  @zobinHuang optimized the performance of
  the batch-invariant vLLM execution path.