Skip to content

[eldritch] attention: avoid sparse indexer host sync - #52

Merged
lukealonso merged 2 commits into
dev/eldritch-enlightenmentfrom
codex/eldritch-b12x-indexer-hostsync-20260626
Jun 28, 2026
Merged

[eldritch] attention: avoid sparse indexer host sync#52
lukealonso merged 2 commits into
dev/eldritch-enlightenmentfrom
codex/eldritch-b12x-indexer-hostsync-20260626

Conversation

@voipmonitor

@voipmonitor voipmonitor commented Jun 26, 2026

Copy link
Copy Markdown

Summary

Remove decode-time CUDA scalar synchronization from sparse-indexer metadata paths without widening the hot GLM/Kimi-style B12X scorer window.

The regression came from this pattern in decode metadata setup:

decode_topk_max_seq_len = int(seq_lens.max().item())

seq_lens is a CUDA tensor in decode, so this host scalar read can synchronize the stream on every generated token. The scalar is needed only as metadata/scheduler input for the sparse top-k/indexer path; exact per-row seq_lens remains available on device for the kernels.

Scope

This PR affects sparse MLA models that actually use DeepseekV32IndexerMetadataBuilder / sparse_attn_indexer.

Covered paths:

  • B12X_MLA_SPARSE with VLLM_USE_B12X_SPARSE_INDEXER=1.
  • Compressed MLA rows (compress_ratio > 1), e.g. DS4-style compressed KV/indexer layouts.
  • Uncompressed sparse-indexer rows (compress_ratio == 1) when they use the B12X sparse indexer, e.g. GLM-style B12X sparse attention.
  • Non-B12X compressed MLA sparse-indexer path, e.g. DS4 Lucifer / FlashInfer CUTLASS compressed path.

Not covered / not relevant:

  • Normal Kimi K2.6/K2.7 launches that use TRITON_MLA. Those do not enter this sparse-indexer metadata path, so this PR neither helps nor hurts them.
  • Non-B12X uncompressed sparse-indexer launches (compress_ratio == 1 with VLLM_USE_B12X_SPARSE_INDEXER=0) still keep the existing fallback seq_lens.max().item(). If we want to support GLM/Kimi through a non-B12X sparse backend later, that should be handled separately.

Implementation

The fix splits behavior by indexer layout and backend:

  • B12X compressed MLA (compress_ratio > 1): use the existing graph-stable active_width_tokens = ceil(max_seq_len / compress_ratio) bound. The B12X scorer already consumes this live-window bound, while exact per-row device seq_lens remains available to the kernel. This removes the sync without changing the scoring contract.
  • B12X uncompressed sparse-indexer rows (compress_ratio == 1): keep the exact scalar because broad active_width measurably slows GLM. Instead of reading the CUDA tensor, compute the exact max from runner-maintained CPU seq-len shadows. For DCP this prefers dcp_local_seq_lens_cpu, so the scalar stays in the same coordinate system as the rank-local seq_lens tensor.
  • Async-spec fallback: if the CPU shadow is not authoritative, fall back to the graph-stable active-width bound instead of synchronizing on CUDA.
  • Non-B12X compressed MLA: use the metadata upper bound for the host scalar. DeepGEMM still receives exact per-row device seq_lens; only the host max_seq_len scorer bound avoids the CUDA reduction.

No B12X workspace/arena behavior is introduced. The vLLM path remains eager and caller-scratch-owned; this PR only changes metadata scalar selection.

Relationship To Upstream PR vllm-project#46178

Upstream vllm-project/vllm#46178 implements the broader DCP sparse-attention algorithm: DCP-local lengths, sparse top-k candidate all-gather, global top-k selection, local/global remap, MTP handling, and FP8 KV support for sparse MLA.

This PR is narrower. It addresses the decode-time host-sync issue in our eldritch/B12X integration and the compressed non-B12X sparse-indexer path. It does not replace the global-top-k algorithm; it only makes the metadata scalar computation sync-free in the paths we serve.

Validation

Static checks:

python3 -m py_compile \
  vllm/v1/attention/backends/mla/indexer.py \
  tests/model_executor/layers/test_sparse_attn_indexer_b12x.py

Focused helper coverage was added for:

  • exact CPU-shadow max
  • DCP-local CPU-shadow max
  • flattened variable-length MTP padded rows
  • missing/non-authoritative CPU-shadow fallback

GLM-5.2 B12X Sparse Indexer

Launch shape:

  • GLM-5.2 Luke NVFP4
  • TP8 / DCP1 / MTP off / A16
  • B12X_MLA_SPARSE
  • GPUs 8-15
  • max_num_seqs=1
  • max_cudagraph_capture_size=4
  • fixed 512-token cc1 decode
Variant 3x 512-token cc1 avg test.py -L gen tok/s CJK
baseline with CUDA .item() 77.56 77.52 0
broad active-width B12X variant 72.43 72.31 0
pure revert of sync-introducing commit 78.50 78.43 0
this PR, exact CPU shadow 78.20 78.09 0

Result: GLM stays close to the pure-revert speed while removing the CUDA scalar sync from the B12X exact path.

DS4 / Non-B12X Compressed Indexer

The second commit extends the no-host-sync rule to the non-B12X compressed MLA path (compress_ratio > 1) used by DS4 Lucifer / FlashInfer CUTLASS.

A/B was run on:

  • same image
  • same GPU pair 8,9
  • DS4 Flash TP2 no-MTP
  • debug startup settings only: max_num_seqs=1, max_cudagraph_capture_size=4, max_num_batched_tokens=2048
Variant 512-token cc1 gen tok/s CUDA sync observation
unpatched non-B12X compressed path 117.77 cudaStreamSynchronize 1022x / 6.88s in nsys runtime stats
this PR with compressed metadata bound 124.66 (124.72-124.76 on reruns) cudaStreamSynchronize no longer appears in runtime top; remaining sync is dominated by the existing event path

Reports:

  • /root/bench-results/ds4-veloq-cc1-v3-vs-current-20260626/current-v420-unpatched-cutlass-cc1-debugseq1-gpu89-20260626-124648.nsys-rep
  • /root/bench-results/ds4-veloq-cc1-v3-vs-current-20260626/current-v420-syncfix-cutlass-cc1-debugseq1-gpu89-20260626-123827.nsys-rep

@coderabbitai

coderabbitai Bot commented Jun 26, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

DeepseekV32IndexerBackend.build now derives decode top-k limits from CPU shadow seq-lens data when available, carries DCP-local CPU seq-lens through decode setup, and builds schedule metadata after selecting the decode top-k length.

Changes

MLA decode metadata ordering

Layer / File(s) Summary
Decode top-k helper
vllm/v1/attention/backends/mla/indexer.py, tests/model_executor/layers/test_sparse_attn_indexer_b12x.py
_decode_topk_max_seq_len_from_cpu derives decode top-k limits from CPU seq-lens shadows, handles flattened MTP padding rows, and returns None when no authoritative CPU shadow is available.
Decode state init
vllm/v1/attention/backends/mla/indexer.py
build() carries dcp_local_seq_lens_cpu through decode setup, initializes active-width tracking earlier, selects B12X decode top-k length before schedule metadata, and packages the final decode metadata.

Estimated review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

  • local-inference-lab/vllm#4: Also changes vllm/v1/attention/backends/mla/indexer.py’s decode metadata construction around DCP/local decode seq-lens handling.
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly matches the main change: removing host synchronization in the sparse attention indexer path.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch codex/eldritch-b12x-indexer-hostsync-20260626

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

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

@voipmonitor
voipmonitor force-pushed the codex/eldritch-b12x-indexer-hostsync-20260626 branch from fa773ce to 7c05b6a Compare June 26, 2026 06:43
@voipmonitor
voipmonitor force-pushed the codex/eldritch-b12x-indexer-hostsync-20260626 branch from 7c05b6a to 3423f72 Compare June 26, 2026 07:42

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/model_executor/layers/test_sparse_attn_indexer_b12x.py`:
- Line 1188: The test module uses np in sparse attention indexer cases but never
imports numpy, so add the missing top-level import in
test_sparse_attn_indexer_b12x.py near the other imports. Update the module used
by the sparse attn tests so the np.array calls in the affected test functions
resolve correctly and do not raise NameError at runtime.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: adbc1593-64f3-4a70-9e38-9352a29844e6

📥 Commits

Reviewing files that changed from the base of the PR and between fa773cec156fa9c72ef71fbe792464b08a8c591d and 3423f72.

📒 Files selected for processing (2)
  • tests/model_executor/layers/test_sparse_attn_indexer_b12x.py
  • vllm/v1/attention/backends/mla/indexer.py

builder._decode_topk_max_seq_len_from_cpu(
common,
num_decodes=3,
decode_lens_np=np.array([1, 1, 1], dtype=np.int32),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
fd -t f 'test_sparse_attn_indexer_b12x.py' tests --exec rg -nP '^\s*(import\s+numpy|import\s+numpy\s+as\s+np|from\s+numpy\s+import)' {}
echo "---- np usages ----"
fd -t f 'test_sparse_attn_indexer_b12x.py' tests --exec rg -nP '\bnp\.' {} | head

Repository: local-inference-lab/vllm

Length of output: 446


numpy is not imported in this test module, which will cause NameError at runtime.

Static analysis and file inspection confirm that np is used at lines 1188, 1221, 1253, and 1285 without a corresponding import numpy as np statement in the module.

Add the following import near the other top-level imports to resolve the NameError:

import numpy as np
Affected lines
1188:             decode_lens_np=np.array([1, 1, 1], dtype=np.int32),
1221:             decode_lens_np=np.array([1, 1], dtype=np.int32),
1253:             decode_lens_np=np.array([3, 0], dtype=np.int32),
1285:             decode_lens_np=np.array([1], dtype=np.int32),
🧰 Tools
🪛 Ruff (0.15.18)

[error] 1188-1188: Undefined name np

(F821)


[error] 1188-1188: Undefined name np

(F821)

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/model_executor/layers/test_sparse_attn_indexer_b12x.py` at line 1188,
The test module uses np in sparse attention indexer cases but never imports
numpy, so add the missing top-level import in test_sparse_attn_indexer_b12x.py
near the other imports. Update the module used by the sparse attn tests so the
np.array calls in the affected test functions resolve correctly and do not raise
NameError at runtime.

Source: Linters/SAST tools

@voipmonitor voipmonitor changed the title [eldritch] attention: avoid B12X sparse indexer host sync [eldritch] attention: avoid sparse indexer host sync Jun 26, 2026
@lukealonso
lukealonso merged commit 4028c95 into dev/eldritch-enlightenment Jun 28, 2026
2 of 3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants