[eldritch] attention: support DCP in SM120 FlashInfer sparse MLA

[voipmonitor]

Summary

This is a cleaned-up SM120 FlashInfer sparse-MLA DCP patch for the Eldritch release branch (codex/eldritch-enlightenment-release-20260627). It makes FLASHINFER_MLA_SPARSE_SM120 usable with --decode-context-parallel-size > 1 without carrying the unrelated B12X hybrid/debug experiments from earlier overlays.

The PR touches only the FlashInfer sparse MLA path plus one generic warmup shape guard needed for DCP production launches:

• keep one FlashInfer sparse-MLA scratch buffer per CUDA device instead of one process-global tensor;
• initialize the SM120 backend's DCP/CP bookkeeping from vLLM parallel config;
• convert global sparse top-k token ids to this DCP rank's local physical KV slots;
• pass per-token valid counts as seq_lens to the FlashInfer SM120 launcher;
• request and return decode softmax LSE so the common MLA DCP merge can reduce partial attention results;
• make mixed prefill+decode warmup use a DCP-valid synthetic decode prompt, so production max_num_seqs > 1 launches do not autotune invalid zero-local-KV DCP shapes.

No B12X binding, arena, workspace, or kernel path is changed here. The B12X sparse MLA implementation was used only as the behavioral reference for DCP index conversion and LSE return semantics.

Why

Before this change, FLASHINFER_MLA_SPARSE_SM120 was effectively DCP1-only. Under DCP, the common MLA attention layer requires each backend to return decode softmax LSE so the partial results from DCP ranks can be merged correctly. The SM120 FlashInfer backend returned (out, None), so DCP serving had to use B12X_MLA_SPARSE even when the SM120 FlashInfer sparse decode kernel was faster.

The SM120 FlashInfer launcher already supports returning LSE. The missing pieces were the vLLM-side DCP index mapping, valid-count plumbing, head/LSE buffer sizing after DCP query gather, and per-device scratch allocation that is visible during memory profiling.

A second issue showed up only in production-style launches. run_mixed_prefill_decode_warmup() used a two-token synthetic decode prompt. With DCP sparse decode, that prompt can occupy only the first cache block and leave later DCP ranks with an empty local KV range. The warmup then autotunes an invalid DCP shape and can corrupt startup. The fix mirrors the existing DCP warmup constraint: span at least one KV block per DCP rank for the synthetic decode request.

Implementation Notes

The global-to-local DCP mapping reuses triton_convert_dcp_global_index_to_local_index, keeping the sparse top-k contract aligned with the existing DCP-capable sparse MLA implementation.

The SM120 path sizes output and LSE buffers from the runtime query head count (q.shape[1]) rather than the constructor's local num_heads, because DCP decode can present a query tensor whose head dimension has already been gathered across DCP ranks.

The FlashInfer scratch buffer is still caller scratch for the FlashInfer launcher, but it is keyed by CUDA device and allocated during backend initialization. This avoids cross-device tensor reuse and avoids first-decode lazy allocation after vLLM has already sized KV cache.

Validation

Validated on 8x RTX PRO 6000 / SM120 with GLM-5.2 NVFP4, TP8/DCP2 and TP8/DCP4, MTP off, --attention-backend FLASHINFER_MLA_SPARSE_SM120, --moe-backend b12x, --kv-cache-dtype fp8, V2 runner.

Static checks:

python3 -m py_compile \
  vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py \
  vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py \
  vllm/v1/worker/gpu/warmup.py

Debug/smoke validation, TP8/DCP4, max_num_seqs=1, graph cap 4:

python3 /mnt/test.py --port 5966 -L
11 iterations before timeout interrupt
8746 completion tokens
CJK: 0
average generation-only throughput: 67.07 tok/s
average incl. TTFT throughput: 65.98 tok/s

Production-start validation, TP8/DCP2, max_num_seqs=32, graph cap 4:

• startup completes with Application startup complete;
• KV cache size: 1,357,696 tokens;
• mixed prefill+decode autotune completes without the previous cublas/NCCL illegal-address cascade;
• short smoke produces coherent Python, CJK 0, generation-only throughput around 67.4 tok/s.

Observed startup/log checks:

• Using V2 Model Runner
• Using AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120 backend
• Using 'B12X' NvFp4 MoE backend
• TP all-reduce uses ['B12X_PCIE_ONESHOT', 'PYNCCL'] for decode-sized tensors; DCP/EP groups remain PyNCCL.

The stored B12X DCP4/no-MTP reference for the same GLM recipe is about 62.3 tok/s, so the SM120 path is consistently faster on the coding smoke while preserving coherent output.

Scope / Limitations

This validates the ag_rs DCP flow used on our PCIe SM120 machines. I did not claim a2a coverage here.

The path assumes the same sparse MLA prerequisites already required by the backend: packed fp8_ds_mla KV cache layout, a FlashInfer build with the SM120 sparse MLA decode API, and a global top-k sparse indexer contract that can be mapped to local DCP KV slots.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• dev/*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 7bbac3b9-a463-4458-970c-ff85775a6faa

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/eldritch-sm120-dcp-clean-pr-20260628

---

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.}

[lllllinux]

Hi, @voipmonitor

I'm trying to enable DCP (DCP size = 8) by following the deployment guide in the ds4dspark-v8 documentation:

My configuration is:

• Backend: lucifer-cutlass
• TP = 8
• Standard MTP = 2

However, the engine hangs during the flashinfer_sparse_mla_warmup stage and eventually hits an NCCL timeout.

If I disable DCP (or set the DCP size to 1), everything works normally.

Is there any recommended configuration or best practice for enabling DCP with DeepSeek V4 Flash when using TP=8? My goal is to achieve high-concurrency serving.

Thanks!

oh,by the way, could you please update the discord invite link in rtx6kpro