[GG] perf(attention): restore B12X MTP decode fast path - #164
Conversation
📝 WalkthroughWalkthroughB12x sparse MLA now tracks speculative decode batches in attention metadata, parses expanded extend-as-decode modes, adjusts speculative query-row planning, and routes kernels using metadata. Tests add scratch-capacity and causality coverage and extend shared validation to FlashInfer SM120. ChangesSparse MLA speculative decode
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AttentionMetadataBuilder
participant B12xMLASparseImpl
participant SparseMLAKernel
AttentionMetadataBuilder->>B12xMLASparseImpl: provide is_spec_decode
B12xMLASparseImpl->>B12xMLASparseImpl: choose decode or extend path
B12xMLASparseImpl->>SparseMLAKernel: execute selected sparse MLA kernel
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/v1/attention/test_sparse_mla_backends.py (1)
1015-1020: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse the
current_platformabstraction instead oftorch.cudadirectly.This test consistently gates other backends via
current_platform.get_device_capability()/current_platform.has_device_capability(...)(e.g. theFlashInferMLASparseTRTLLMBackendcheck a few lines above). This new line callstorch.cuda.get_device_capability()directly, unconditionally, for every backend in the parametrization — inconsistent with the surrounding pattern and less portable if this shared test is ever exercised on a non-CUDA accelerator path.♻️ Suggested fix
- uses_pow2_scales = torch.cuda.get_device_capability()[ - 0 - ] >= 10 and backend_cls not in ( + uses_pow2_scales = current_platform.get_device_capability().major >= 10 and ( + backend_cls not in ( B12xMLASparseBackend, FlashInferMLASparseSM120Backend, - ) + ) + )🤖 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/v1/attention/test_sparse_mla_backends.py` around lines 1015 - 1020, Replace the direct torch.cuda.get_device_capability() call in the uses_pow2_scales condition with the established current_platform.get_device_capability() abstraction. Preserve the existing capability threshold and backend exclusions, matching the surrounding backend checks.
🤖 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 `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py`:
- Around line 666-673: Update the spec-decode detection in the surrounding
attention setup to avoid converting CUDA-backed cm.is_prefilling through
torch.any and bool. Determine the value from CPU-side scheduler metadata, or
defer the decision until kernel dispatch, while preserving the existing
speculative-token and query-length conditions.
---
Nitpick comments:
In `@tests/v1/attention/test_sparse_mla_backends.py`:
- Around line 1015-1020: Replace the direct torch.cuda.get_device_capability()
call in the uses_pow2_scales condition with the established
current_platform.get_device_capability() abstraction. Preserve the existing
capability threshold and backend exclusions, matching the surrounding backend
checks.
🪄 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: 86def0a5-5a4e-4149-82cf-3f6a29448fb6
📒 Files selected for processing (2)
tests/v1/attention/test_sparse_mla_backends.pyvllm/v1/attention/backends/mla/b12x_mla_sparse.py
| is_spec_decode = False | ||
| if ( | ||
| self.num_speculative_tokens > 0 | ||
| and 1 < cm.max_query_len <= self.num_speculative_tokens + 1 | ||
| and cm.is_prefilling is not None | ||
| ): | ||
| is_spec_decode = not bool(torch.any(cm.is_prefilling[: cm.num_reqs])) | ||
|
|
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate CommonAttentionMetadata's is_prefilling field and its population site.
rg -n -B2 -A2 'is_prefilling' --type=py -g '!**/test*' | head -100
ast-grep run --pattern 'class CommonAttentionMetadata { $$$ }' --lang pythonRepository: local-inference-lab/vllm
Length of output: 2070
🌐 Web query:
PyTorch bool(torch.any(cuda_tensor)) host device synchronization docs item() CUDA tensor
💡 Result:
Using bool on a CUDA tensor, such as bool(torch.any(cuda_tensor)), triggers a host-device synchronization [1][2]. This occurs because the CPU must retrieve the result from the GPU to determine the truth value needed for Python's control flow, which forces the GPU to complete all pending operations in the relevant stream before the value can be returned to the CPU [1][2]. Host-Device Synchronization Mechanics: Operations that require the host to know the value of data residing on the device—such as using .item(), passing a tensor to a Python if-statement, or evaluating bool()—are inherently blocking [1][2]. The CPU cannot proceed with the Python-level conditional logic until the computation on the GPU is finished and the result is transferred back [1][2]. This behavior is a common cause of performance bottlenecks and is incompatible with CUDA graph capture, which requires synchronization-free execution [1][3][4]. Recommendations and Best Practices: To avoid or mitigate these synchronization points, consider the following strategies: 1. Prefer GPU-Native Operations: Use PyTorch operators that perform reductions or conditional logic entirely on the GPU, such as torch.where(), torch.max(), or torch.any(tensor). By keeping the data on the device, you avoid the need for the CPU to materialize the result [1]. 2. Decouple Control Flow: If possible, structure your code to avoid data-dependent branching. For example, instead of using an if-statement that branches based on a tensor's value, compute the operations for both paths (or use masks) and combine them using torch.where() [1]. 3. Async Pipelines: If you must check for completion or a condition (e.g., in a training loop or while replaying CUDA graphs), use an asynchronous pipeline with pinned memory (pin_memory()) and a separate CUDA stream [3]. By copying the result to a CPU tensor asynchronously (non_blocking=True) and checking that CPU-side memory later, you can avoid stalling the main GPU compute stream [5][3]. Note: While PyTorch has made internal improvements to avoid unnecessary synchronization (e.g., for single-element reductions in some cases) [6], any operation that fundamentally requires the host to read a value from a CUDA tensor will remain a synchronization point [1][2]. Always use profiling tools (e.g., torch.cuda.set_sync_debug_mode("warn")) to identify unexpected synchronization in your code [2].
Citations:
- 1: https://docs.nvidia.com/dl-cuda-graph/latest/torch-cuda-graph/sync-free-code.html
- 2: https://discuss.pytorch.org/t/slow-iteration-over-tensor-elements-slow-any/191998
- 3: https://docs.nvidia.com/dl-cuda-graph/examples/rnnt.html
- 4: https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/model_executor/models/indextts2/dit_cuda_graph/
- 5: https://docs.pytorch.org/tutorials/intermediate/pinmem%5Fnonblock.html
- 6: OSU-Nowlab/pytorch@0b363c5
🌐 Web query:
vLLM CommonAttentionMetadata is_prefilling device placement
💡 Result:
In vLLM's architecture, CommonAttentionMetadata is a dataclass used for storing per-batch attention information that is shared across different layers and backends [1][2]. The is_prefilling field within this class is a (batch_size,) boolean tensor used by various attention backends to distinguish between requests currently in the prefill phase and those performing decodes [1]. Regarding device placement and management: 1. Device Placement: CommonAttentionMetadata tensors are typically stored on the GPU to facilitate direct access during kernel execution [3]. However, for certain fields, vLLM maintains both GPU and CPU versions to optimize performance and prevent unnecessary host-to-device (H2D) transfers [1][2][4]. 2. is_prefilling Computation: This field is computed as a comparison: is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu [5]. Because it relies on values like num_computed_tokens and num_prompt_tokens, which are often tracked on the CPU for scheduler logic, it is essential that these tensors are properly handled—especially in scenarios involving padded rows [5]. 3. Padding Issues: A known technical challenge involves stale metadata in padded rows. If padded rows in the CPU request-state tensors are not properly cleared or neutralized, the is_prefilling computation can incorrectly classify an inactive padded row as a prefill row (e.g., in Mamba backends), potentially causing errors in CUDA graph replays [5]. To address this, vLLM implementations generally ensure that metadata for padded rows is cloned, zeroed, or otherwise neutralized before being utilized in attention calculations [5]. In summary, is_prefilling is a critical flag for distinguishing execution phases, and its accuracy relies on managing CPU-side request-state tensors and correctly handling padding before it is used by attention backends [1][5].
Citations:
- 1: https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backend.py
- 2: https://github.com/Deepfocused/vllm-exaone4.0/blob/main/vllm/v1/attention/backends/utils.py
- 3: https://docs.vllm.ai/en/v0.10.1/api/vllm/attention/backends/utils.html
- 4: [Attention] Refactor attention metadata builder interface vllm-project/vllm#20466
- 5: Stale padded request metadata can misclassify Mamba CUDA graph rows vllm-project/vllm#41841
🌐 Web query:
site:github.com/vllm-project/vllm "is_prefilling =" "CommonAttentionMetadata"
💡 Result:
In the vLLM codebase, is_prefilling is a field within the CommonAttentionMetadata class [1][2][3]. It is defined as a torch.Tensor (specifically a boolean tensor of shape (batch_size,)) that indicates whether a given request is still in the prefill phase, defined as the condition where the number of computed tokens is less than the number of prompt tokens (num_computed_tokens < num_prompt_tokens) [1][2][4]. This field is utilized by various attention backends to distinguish between actual decode requests and "short extends" (chunked prefills) that might otherwise fall under a decode threshold [1][5][6]. By checking this flag, the system can correctly categorize request types—decode, short_extend, long_extend, or prefill—allowing for more accurate batch reordering and processing, particularly in attention kernels that treat these categories differently [5][6].
Citations:
- 1: vllm-project@e1d85e5
- 2: https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backend.py
- 3: https://github.com/vllm-project/vllm/blob/17d87168/vllm/v1/attention/backend.py
- 4: https://github.com/vllm-project/vllm/blob/7cc302dd/vllm/v1/attention/backend.py
- 5: [Attention] Support distinguishing between short extends and decodes vllm-project/vllm#37303
- 6: https://github.com/vllm-project/vllm/blob/7c2acd38/vllm/v1/attention/backends/utils.py
🌐 Web query:
site:github.com/vllm-project/vllm "cm.is_prefilling" "vllm/v1/attention/backends/mla/b12x_mla_sparse.py"
💡 Result:
In the vLLM project, specifically within the v1 architecture, cm.is_prefilling is a property or flag used to determine if the engine or a specific attention operation is currently in the prefill (prompt processing) phase [1][2]. Regarding the file path vllm/v1/attention/backends/mla/b12x_mla_sparse.py, this refers to a component within vLLM's attention backends, likely introduced to support Sparse Multi-Head Latent Attention (MLA) for specific hardware architectures (e.g., SM80, SM12x/Blackwell) [3][4]. Sparse MLA support in vLLM v1 involves multiple backends, such as FLASHMLA_SPARSE, FLASH_ATTN_MLA_SPARSE, and TRITON_MLA_SPARSE, which are used depending on hardware compatibility, KV cache data types, and performance requirements [3][5][4][6]. The b12x identifier typically relates to Blackwell or specialized SM12x-series hardware support [3][4]. The cm object in cm.is_prefilling often refers to a context manager or a configuration/metadata object passed into attention kernels or scheduling modules to toggle behavior between prefill and decode (or between fixed/variable sequence lengths) [1][2][7]. For example, when chunked prefill is enabled, this flag helps the scheduler or kernel distinguish between stages to optimize memory access and kernel dispatch, such as opting for eager execution for variable shapes during verification phases versus CUDA graphs for fixed shapes [1][2][7].
Citations:
- 1: [Core] Don't schedule spec tokens with prefill chunks vllm-project/vllm#33652
- 2: [Chunked Prefill][4/n] Chunked prefill scheduler. vllm-project/vllm#3853
- 3: [Feature] TRITON_MLA_SPARSE backend for SM8x/11x/12x DSA Sparse MLA Support vllm-project/vllm#38476
- 4: [Attention] TRITON_MLA_SPARSE backend for SM80/SM121 sparse MLA (rebase & takeover of #38476) vllm-project/vllm#47629
- 5: https://github.com/vllm-project/vllm/blob/7fe7fa9c/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
- 6: [Attention] Add FLASH_ATTN_MLA_SPARSE backend for Hopper sparse MLA vllm-project/vllm#46189
- 7: [Bug]: b12x NSA+MTP speculative decoding hangs on PCIe TP=8 — NCCL topology-aware scheduling fix vllm-project/vllm#43315
🌐 Web query:
site:github.com/vllm-project/vllm "is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu" "backend.py"
💡 Result:
The expression is_prefilling = num_computed_tokens_cpu < num_prompt_tokens_cpu is a logic check used within the vLLM engine (specifically within the V1 attention backend) to determine if a request is in the prefill phase [1][2]. In this context: - num_computed_tokens_cpu tracks the number of tokens already processed for a request [1][2]. - num_prompt_tokens_cpu tracks the total number of tokens in the prompt [1][2]. - When the number of computed tokens is strictly less than the total prompt tokens, the request is identified as being in the prefill phase (is_prefilling = True) [1][3]. This logic helps the engine distinguish between prefill (processing prompt tokens) and decode (generating new tokens) phases, allowing for specialized scheduling and kernel selection [3][4]. However, this specific implementation has been subject to various refinements, including: - Deprecation concerns: Usage of num_computed_tokens_cpu has been marked for deprecation in favor of more robust state tracking to improve compatibility with asynchronous scheduling [5][2]. - Padding issues: Stale data in padded batch rows can cause inactive requests to be incorrectly flagged as prefills, which requires careful cleaning or cloning/zeroing of these fields before the comparison is performed [1]. - Mamba-specific adjustments: For certain models like Mamba, the standard prefill/decode classification is sometimes insufficient. For instance, single-token "extends" (where a prefill process continues) may be forced to behave as decodes for performance reasons, requiring the logic to be manually overridden or ORed with other flags [3][6].
Citations:
- 1: Stale padded request metadata can misclassify Mamba CUDA graph rows vllm-project/vllm#41841
- 2: [Attention][1/n] Remove usage of deprecated
seq_lens_cpuandnum_computed_tokens_cpuCommonAttentionMetadata properties vllm-project/vllm#31773 - 3: [Bugfix] mamba: run single-token extends as decodes vllm-project/vllm#42430
- 4: https://github.com/vllm-project/vllm/blob/469f3dcf/vllm/v1/attention/backends/utils.py
- 5: [Bugfix] Fix MTP edge case in split_decodes_and_prefills vllm-project/vllm#32716
- 6: https://github.com/vllm-project/vllm/blob/3fd9d2d3/vllm/v1/attention/backends/mamba_attn.py
🌐 Web query:
site:github.com/vllm-project/vllm "class CommonAttentionMetadata" "is_prefilling: torch.Tensor"
💡 Result:
In the vLLM project, CommonAttentionMetadata is a dataclass used to store per-batch attention metadata shared across different layers and backends [1]. The field is_prefilling is defined as a torch.Tensor with a shape of (batch_size,) [1]. It contains boolean values indicating whether a specific request is still in the prefill phase, determined by the condition where the number of computed tokens is less than the number of prompt tokens [1][2]. This tensor is primarily used by various attention backends to distinguish between actual decode operations and short extensions (where a request has a small query length but is still technically in the prefill phase) [1][2]. When splitting batches into decodes and prefills, this field can be used to ensure that short extends are correctly categorized if the system configuration does not treat them as decodes [3][2]. Because this tensor is part of the request-state metadata, care must be taken in scenarios like CUDA graph replaying to ensure that padded or inactive rows do not contain stale values, which could otherwise lead to incorrect classification of requests [4].
Citations:
- 1: https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backend.py
- 2: vllm-project@e1d85e5
- 3: https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/utils.py
- 4: Stale padded request metadata can misclassify Mamba CUDA graph rows vllm-project/vllm#41841
Avoid the host sync on the spec-decode path. bool(torch.any(cm.is_prefilling[: cm.num_reqs])) will block if cm.is_prefilling is CUDA-backed. Keep this decision on CPU-side scheduler metadata, or defer it until kernel dispatch.
🤖 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 `@vllm/v1/attention/backends/mla/b12x_mla_sparse.py` around lines 666 - 673,
Update the spec-decode detection in the surrounding attention setup to avoid
converting CUDA-backed cm.is_prefilling through torch.any and bool. Determine
the value from CPU-side scheduler metadata, or defer the decision until kernel
dispatch, while preserving the existing speculative-token and query-length
conditions.
|
Combined-stack validation exposed five pre-existing sparse-MLA tests that need their |
2 participants
Summary
0as a kill switch and1as an explicit legacy force modeRoot cause
VLLM_B12X_MLA_SPEC_EXTEND_AS_DECODEdefaulted to disabled because the decode kernel was assumed to treat verifier rows as independent one-token queries. That assumption does not match the current path: candidate KV rows are already present and B12X receives a causal cache length for every verifier token, so each later token can attend to earlier candidates from the same verification step.As a result, MTP verification unnecessarily used the substantially slower sparse extend kernel. This accounted for the observed v20 B12X MTP decode gap against FlashInfer.
Behavior
The setting now accepts three modes:
auto(default): use decode only for a genuine speculative-verification batch0: always keep multi-token batches on extend1: force the old opt-in behavior for every eligible short extendAuto classification requires a configured speculative width,
1 < max_query_len <= 1 + num_speculative_tokens, and no real request still in prompt prefill. Short chunked prefills therefore stay on the extend path.Scratch capacity is exact in the normal modes: one row/request when disabled and
1 + num_speculative_tokensrows/request in auto mode. Only explicit force mode reserves up toVLLM_B12X_MLA_SPEC_DECODE_MAX_Q.Validation
SM120 targeted tests on the v20 CUDA 13.2 stack: 10 passed.
ruff check,ruff format --check, andgit diff --checkAdditional full-model differential validation over 16 verifier rows:
0.0014908350.0000015070.0049368910.0000000320.01994660780-100%The distribution is strongly skewed: most rows are effectively identical, while
the maximum came from the fourth position of one MTP3 verifier batch. Later
candidate positions generally show more floating-point divergence because they
include more candidate KV rows. Across 632 layer/rank comparisons, maximum
absolute error was
0.0078125and maximum relative L2 error was0.00057.This is not bitwise equivalence between the sparse extend and decode kernels,
but it rules out the original causality concern: both paths independently match
token-wise causal SDPA, and no tested full-model row changed its top-1 token.
A conventional checkpoint-versus-BF16 KLD run would not exercise this branch,
because it does not create a multi-token MTP target-verification batch. A larger
paired verifier-logit corpus would still be required to claim zero statistical
quality impact beyond this targeted validation.
DCP2, DCP4, and DCP8 MTP verification completed coherently with no CUDA/Xid fault. Acceptance was
66/93,66/93, and67/90, respectively.Controlled performance
GLM-5.2 NVFP4, TP8/DCP1/MTP3, A16 + online MXFP8, CC1. B12X runs used the same v20 image and launch parameters; only this module was overlaid.
The optimized B12X path is about 6-7% faster than the locally reproduced FlashInfer ctx0 result, 4.6% faster at 64k context, and retains an 18.3% prefill advantage.