The Virtual Block Table: Reusing Triton MLA for Sparse Attention on Blackwell GPUs

Introduction

In the course of deploying the GLM-5 model on a machine equipped with NVIDIA RTX PRO 6000 Blackwell GPUs, a team of engineers encountered a seemingly insurmountable blocker: no existing attention backend in the vLLM inference framework supported the combination of Blackwell's SM120 compute capability, sparse Multi-Head Latent Attention (MLA) with a DSA indexer, and a qk_nope_head_dim of 192. After days of debugging GGUF loader patches, fixing weight corruption bugs, and wrestling with CUDA toolkit compatibility, the final wall was the attention backend itself. Message [msg 1739] captures the moment of breakthrough—the assistant's realization that the existing Triton MLA decode kernel could be reused for sparse attention with almost no modification, by treating the sparse physical indices as a "virtual block table" with page_size=1. This insight, elegant in its simplicity, enabled the team to bypass months of CUDA kernel development and get the 402GB GLM-5 model loading onto the Blackwell GPUs.

The Context: A Cascade of Blockers

To understand the significance of message [msg 1739], we must first appreciate the gauntlet of failures that preceded it. The team had been working for days to deploy the GLM-5 model—a massive sparse Mixture-of-Experts architecture using DeepSeek-style Attention (DSA) with a sparse MLA indexer. The deployment target was a machine with eight RTX PRO 6000 Blackwell GPUs (compute capability SM120), running Ubuntu 24.04 with CUDA 13.1.

The journey had already been arduous. The team had patched vLLM's gguf_loader.py to support the glm_moe_dsa architecture, fixed a latent bug in DeepSeek V2/V3 GGUF support where kv_b_proj tensors were being incorrectly mapped, built llama-gguf-split from source to merge ten split GGUF files into a single 402GB file, and revised the KV cache reassembly logic after discovering that the tensors used an older shape representation.

When the first vllm serve launch was attempted, three distinct failures emerged in rapid succession. First, the maybe_override_with_speculators function crashed on the unsupported glm-dsa architecture—fixed by patching the model config. Second, a torch.bfloat16 dtype incompatibility with GGUF quantization required adding --dtype float16 to the launch command. Third, and most critically, there was no valid attention backend for the Blackwell SM120 GPU that could handle sparse MLA with the model's qk_nope_head_dim=192.

The assistant had investigated whether updating to the latest vLLM nightly would help, but discovered that the newest FlashInferMLA sparse backend (added in PR #33451 on February 12, 2026) only supported SM100 (not SM120) and required qk_nope_head_dim=128 (GLM-5 has 192). The existing Triton MLA backend supported all compute capabilities but had no sparse attention support. The fundamental problem was clear: no existing MLA backend in vLLM supported the combination of SM120 + sparse attention + qk_nope_head_dim=192. Updating vLLM would not solve this—the support simply did not exist upstream.

The user and assistant agreed on a strategy: patch the Triton MLA backend to support sparse attention. The Triton backend is pure Python/Triton code, not a compiled CUDA kernel, making it the most feasible target for modification ([msg 1725]).

The Investigation: Understanding the Sparse MLA Architecture

Before the breakthrough in message [msg 1739], the assistant conducted a thorough investigation of the sparse MLA architecture. Using a subagent task ([msg 1727]), it read the full contents of three critical files on the remote machine: the SparseMLAAttentionImpl base class, the FlashMLASparseBackend implementation, and the TritonMLABackend implementation. The subagent returned a comprehensive analysis of the class hierarchy, showing that SparseMLAAttentionImpl inherits from MLAAttentionImpl and overrides the attention metadata and forward methods to handle sparse indices.

The assistant then examined the Triton decode kernel itself ([msg 1728]-[msg 1731]), discovering that the core kernel _fwd_kernel_stage1 already uses a Req_to_tokens block table to gather KV cache entries by page number. The kernel iterates over sequence positions in blocks of size BLOCK_N, resolving each position to a physical cache slot via Req_to_tokens[req, pos // PAGE_SIZE] * PAGE_SIZE + pos % PAGE_SIZE.

The FlashMLA sparse backend's approach was also studied ([msg 1732]-[msg 1737]). The critical insight came from the _forward_bf16_kv method, which calls triton_convert_req_index_to_global_index to convert per-request sparse indices into global physical cache slots. These physical indices are then passed to the FlashMLA kernel as a pseudo-block-table where each entry is a direct physical cache slot (effectively page_size=1).

By message [msg 1738], the assistant had confirmed the pattern: the FlashMLA sparse backend treats topk_indices as a block table with page_size=1, where each index has already been resolved to a physical cache position.## The Breakthrough: Message [msg 1739]

It was at this point that the assistant had its epiphany. The subject message begins:

Now I understand the full API. The Triton kernel takes: - req_to_token: the block table [num_reqs, max_blocks] — used to resolve page numbers - b_seq_len: sequence lengths per request - page_size: the block/page size

>

For sparse, the trick is: we convert the topk_indices to physical positions using triton_convert_req_index_to_global_index, then treat them as a virtual block table with page_size=1. The "seq_len" becomes the topk count.

>

This means we can reuse the existing Triton decode kernel almost as-is! We just need: 1. req_to_token = topk_indices_physical reshaped to [num_tokens, topk] 2. b_seq_len = number of valid topk entries per token 3. page_size = 1 4. k_buffer and v_buffer = flat view of the KV cache

The reasoning here is a beautiful example of interface abstraction. The assistant recognized that the Triton decode kernel's API was more general than its typical usage. The req_to_token parameter is just a lookup table mapping logical positions to physical cache slots. In dense attention, the logical positions are contiguous sequence indices [0, seq_len), and the block table maps page numbers to physical page slots. But the kernel doesn't care about the meaning of the indices—it just needs to resolve each position it iterates over to a physical cache address.

For sparse attention, the logical positions are the top-k sparse indices selected by the DSA indexer. After conversion via triton_convert_req_index_to_global_index, these indices become direct physical cache slot numbers. By reshaping them into a [num_tokens, topk] tensor and passing it as req_to_token, the existing kernel can iterate over these positions exactly as it would iterate over dense sequence positions—except now page_size=1 means each "page" contains exactly one element, and the "sequence length" is the number of top-k entries per token.

This insight meant that the assistant could write a TritonMLASparseBackend class that:

  1. Inherits from SparseMLAAttentionImpl (the base class for sparse MLA attention)
  2. In its forward_mqa method, converts sparse indices to physical positions using the existing triton_convert_req_index_to_global_index utility
  3. Calls the existing decode_attention_fwd kernel with the physical indices as the block table, page_size=1, and the top-k count as the sequence length No new CUDA kernels needed. No Triton kernel modifications. Just a clever repurposing of existing interfaces.

Assumptions and Reasoning

The assistant made several key assumptions in this message, all of which turned out to be correct:

Assumption 1: The Triton decode kernel treats req_to_token as a flat lookup table without enforcing any block structure. This was a safe assumption given the kernel's implementation—the kernel uses req_to_token[req, pos // PAGE_SIZE] to get a page number, then multiplies by PAGE_SIZE and adds the offset. With PAGE_SIZE=1, pos // 1 = pos, so the kernel effectively does req_to_token[req, pos] * 1 + pos % 1 = req_to_token[req, pos]. The kernel doesn't care whether the indices represent page numbers or direct physical slots—it just uses them to compute a memory address.

Assumption 2: The k_buffer and v_buffer can be viewed as flat tensors. The KV cache in vLLM is stored as a flat buffer indexed by physical slot number. The Triton kernel already accesses it via k_buffer[physical_slot, head, dim], so passing a flat view of the KV cache works identically for sparse and dense attention.

Assumption 3: The top-k indices, after conversion to physical positions, have the same shape semantics as a block table. The FlashMLA sparse backend already proved this pattern worked—it passed the converted indices to flash_mla_sparse_fwd as a pseudo-block-table. The assistant was essentially adopting the same approach for the Triton backend.

Assumption 4: The number of valid top-k entries per token can serve as b_seq_len. In sparse attention, each token may have a different number of valid sparse indices (though typically all tokens in a batch have the same top-k count). The kernel uses b_seq_len to determine how many positions to iterate over, so passing the top-k count works naturally.

The Thinking Process

The assistant's thinking process in this message is remarkably clear and structured. It begins by summarizing the Triton kernel's API in its own words, confirming understanding of the three key parameters (req_to_token, b_seq_len, page_size). Then it states the core insight as a "trick"—converting topk_indices to physical positions and treating them as a virtual block table.

The thinking then walks through the implications step by step. Each of the four numbered points maps a sparse attention concept to a Triton kernel parameter. The mapping is precise: the sparse indices become the block table, the top-k count becomes the sequence length, page_size becomes 1, and the KV cache remains unchanged.

What's particularly notable is the confidence in the conclusion: "This means we can reuse the existing Triton decode kernel almost as-is!" The word "almost" is important—the assistant acknowledges that a new backend class is needed to orchestrate the conversion and calling, but the kernel itself requires zero modification. This is the hallmark of a deep understanding of the system: recognizing that an interface designed for one purpose (dense attention with paged KV cache) is actually general enough to serve another (sparse attention with direct indices).

The message concludes with the assistant immediately acting on the insight, writing the full triton_mla_sparse.py backend file. The LSP errors shown at the end are from unrelated files (decode_latency_breakdown.py and decode_gap_analysis.py), confirming that the new backend file itself had no syntax or type errors.

Input Knowledge Required

To fully understand this message, one needs:

  1. Understanding of vLLM's attention backend architecture: The class hierarchy (AttentionImplBase, MLAAttentionImpl, SparseMLAAttentionImpl), the metadata types, and how backends are registered and selected.
  2. Knowledge of the Triton decode kernel's API: The decode_attention_fwd function signature, the role of req_to_token as a block table, and how page_size and b_seq_len control iteration.
  3. Understanding of sparse MLA attention: The DSA indexer produces top-k sparse indices per token, which select a subset of KV cache positions rather than attending to the full sequence.
  4. Knowledge of the triton_convert_req_index_to_global_index utility: This function converts per-request sparse indices (which are relative to the request's own cache allocation) into global physical cache slot numbers (which are absolute positions in the flat KV cache buffer).
  5. Familiarity with the GLM-5 model architecture: Specifically that it uses DeepSeek-style Attention (DSA) with a sparse MLA indexer and qk_nope_head_dim=192.
  6. Context of the deployment environment: Blackwell SM120 GPUs, CUDA 13.1, the specific vLLM nightly build being used, and the fact that no upstream backend supported the required combination of features.

Output Knowledge Created

This message produced:

  1. A design for the TritonMLASparseBackend: A new attention backend that inherits from SparseMLAAttentionImpl and reuses the existing Triton decode kernel with sparse physical indices as a virtual block table.
  2. The actual implementation file: triton_mla_sparse.py was written to disk, containing the full backend class with the sparse forward logic.
  3. A reusable pattern for sparse attention on Triton: The insight that any set of physical cache positions can be treated as a block table with page_size=1 is general—it could be applied to other sparse attention patterns beyond the DSA indexer used by GLM-5.
  4. Confidence that the approach would work: The assistant's clear reasoning and the fact that the FlashMLA sparse backend already proved the same pattern gave the team confidence to proceed with deployment.

Significance and Impact

The impact of this message cannot be overstated. Without this insight, the team would have faced two equally unpalatable options: write a custom CUDA kernel for sparse MLA on SM120 (months of work), or abandon the Blackwell GPUs entirely and deploy on older hardware. Instead, the assistant recognized that the existing Triton kernel was already capable—it just needed to be fed the right data.

This is a textbook example of leveraging existing abstractions rather than building new ones. The block table abstraction in vLLM's attention kernels was designed for paged KV cache management, but it turned out to be isomorphic to the sparse index table needed for DSA attention. By mapping one concept onto the other, the assistant saved weeks or months of engineering effort.

The message also demonstrates the power of reading the code carefully. The assistant didn't just look at the Triton kernel's function signature—it traced through the kernel's actual memory access pattern, understood how req_to_token was used to resolve positions, and recognized that the resolution logic was independent of whether the positions were contiguous or sparse. This level of detailed understanding is what enabled the elegant solution.

In the broader context of the session, this message marks the turning point. After days of patching loaders, fixing bugs, and debugging launch failures, the attention backend was the final major blocker. Once the TritonMLASparseBackend was implemented and deployed, the model loading progressed past the attention backend selection, successfully selecting TRITON_MLA_SPARSE and beginning to load the 402GB model onto the GPUs. The team had crossed the finish line.