[eldritch] attention: avoid sparse indexer host sync - #52
Conversation
📝 WalkthroughWalkthrough
ChangesMLA decode metadata ordering
Estimated review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
fa773ce to
7c05b6a
Compare
7c05b6a to
3423f72
Compare
There was a problem hiding this comment.
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.pyvllm/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), |
There was a problem hiding this comment.
🎯 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\.' {} | headRepository: 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 npAffected 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
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:
seq_lensis 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-rowseq_lensremains available on device for the kernels.Scope
This PR affects sparse MLA models that actually use
DeepseekV32IndexerMetadataBuilder/sparse_attn_indexer.Covered paths:
B12X_MLA_SPARSEwithVLLM_USE_B12X_SPARSE_INDEXER=1.compress_ratio > 1), e.g. DS4-style compressed KV/indexer layouts.compress_ratio == 1) when they use the B12X sparse indexer, e.g. GLM-style B12X sparse attention.Not covered / not relevant:
TRITON_MLA. Those do not enter this sparse-indexer metadata path, so this PR neither helps nor hurts them.compress_ratio == 1withVLLM_USE_B12X_SPARSE_INDEXER=0) still keep the existing fallbackseq_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:
compress_ratio > 1): use the existing graph-stableactive_width_tokens = ceil(max_seq_len / compress_ratio)bound. The B12X scorer already consumes this live-window bound, while exact per-row deviceseq_lensremains available to the kernel. This removes the sync without changing the scoring contract.compress_ratio == 1): keep the exact scalar because broadactive_widthmeasurably slows GLM. Instead of reading the CUDA tensor, compute the exact max from runner-maintained CPU seq-len shadows. For DCP this prefersdcp_local_seq_lens_cpu, so the scalar stays in the same coordinate system as the rank-localseq_lenstensor.seq_lens; only the hostmax_seq_lenscorer 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#46178implements 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:
Focused helper coverage was added for:
GLM-5.2 B12X Sparse Indexer
Launch shape:
B12X_MLA_SPARSE8-15max_num_seqs=1max_cudagraph_capture_size=4test.py -Lgen tok/s.item()77.5677.52072.4372.31078.5078.43078.2078.090Result: 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:
8,9max_num_seqs=1,max_cudagraph_capture_size=4,max_num_batched_tokens=2048117.77cudaStreamSynchronize1022x / 6.88s in nsys runtime stats124.66(124.72-124.76on reruns)cudaStreamSynchronizeno longer appears in runtime top; remaining sync is dominated by the existing event pathReports:
/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