The Prefill Path: Tracing a Garbage-Output Bug Through vLLM's MLA Attention Backend

Introduction

In the long arc of debugging a complex ML deployment — loading a 402GB GGUF-quantized GLM-5 model onto 8× NVIDIA Blackwell GPUs using vLLM — there comes a moment when the obvious suspects have been eliminated and the investigation must turn to deeper, more architectural questions. Message [msg 1928] captures exactly such a moment. After confirming that the GGUF dequantization kernel produces correct results on SM120 (Blackwell's compute capability), after verifying that the weight name mapping between GGUF tensor names and HuggingFace parameter names works correctly, and after ruling out several other potential failure modes, the assistant is left with a model that loads without errors but generates incoherent output — garbage tokens with flat log-probability distributions. Message 1928 represents a critical pivot point in the investigation: the assistant shifts focus from whether the weights are loaded correctly to whether the attention computation itself works on the Blackwell architecture.

The Message in Full

The subject message contains two distinct parts: a reasoning statement and a diagnostic command. The reasoning reads:

The prefill path uses self.kv_b_proj(kv_c_normed) which is the same ColumnParallelLinear that receives our dequantized weight. And the _run_prefill_new_tokens calls the attention kernel.

The assistant then executes a bash command to inspect the available prefill backends:

ssh root@10.1.230.174 'grep -n "_run_prefill_new_tokens\|_flash_attn_varlen" /root/ml-env/lib/python3.12/site-packages/vllm/model_executor/layers/attention/mla_attention.py | head -10' 2>&1

Which returns:

1993:            self._run_prefill_new_tokens = self._run_prefill_new_tokens_trtllm_ragged
1998:            self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi
2003:            self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn
2014:            self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa
2044:    def _flash_attn_varlen_diff_headdims(
2081:    def _run_prefill_new_tokens_fa(
2084:        return self._flash_attn_varlen_diff_headdims(
2097:    def _run_p...

Why This Message Was Written

The message emerges from a specific logical chain. The assistant has already established that:

The GGUF dequantization kernel works correctly on SM120 (verified with a max diff of 0.00012 against float16 reference).
The weight name mapping from HF parameter names to GGUF tensor names is correct — the gguf-py library's get_tensor_name_map for the glm-dsa architecture correctly maps model.layers.0.self_attn.q_a_proj to blk.0.attn_q_a, model.layers.0.mlp.down_proj to blk.0.ffn_down, and so on.
The model loads without errors and the server starts serving requests. Yet the output is garbage. The logprobs for obvious continuations (predicting "2" after "1 2 3 4 5 6 7 8 9 10") are around -20 to -24, far from the near-zero values a correctly loaded model would produce. The generated tokens are random with logprobs around -4. The assistant's reasoning in this message connects two observations. First, the prefill path (processing the input prompt) calls self.kv_b_proj(kv_c_normed) — this is the same ColumnParallelLinear layer that the assistant had been scrutinizing for a potential tensor-parallelism sharding mismatch. Earlier investigation revealed that the kv_b_proj weight is reassembled from separate k_b and v_b tensors into a full [28672, 512] tensor, but the ColumnParallelLinear expects a TP-sharded [3584, 512] parameter (where 3584 = 28672 / 8 GPUs). Second, the prefill path then calls the attention kernel through _run_prefill_new_tokens. If the attention kernel itself has issues on SM120, that could explain the garbage output independently of the weight loading. This is a classic debugging fork: the assistant must determine whether the bug is in the data (weights not being correctly loaded or sharded) or in the computation (the attention kernel producing incorrect results on Blackwell hardware). Message 1928 represents the decision to investigate the computation path.## The Thinking Process Visible in the Reasoning The assistant's reasoning in this message is remarkably compressed but reveals a clear logical structure. The first sentence — "The prefill path uses self.kv_b_proj(kv_c_normed) which is the same ColumnParallelLinear that receives our dequantized weight" — does two things simultaneously. It acknowledges the ongoing concern about the kv_b_proj weight sharding (the ColumnParallelLinear that the assistant had been debugging in previous messages) while also broadening the investigation to consider the full computational pipeline. The phrase "the same ColumnParallelLinear" is a subtle anchor point: it tells the reader (and the assistant itself) that even if the attention kernel investigation reveals nothing, the kv_b_proj hypothesis remains live. The second sentence — "And the _run_prefill_new_tokens calls the attention kernel" — completes the logical bridge. The prefill path involves both the kv_b_proj linear transformation and the attention kernel. If either is broken, the output is garbage. The assistant is systematically enumerating the components in the forward pass and checking each one. The command that follows is a direct consequence of this reasoning. The assistant searches for _run_prefill_new_tokens and _flash_attn_varlen in the MLA attention implementation to understand which attention backend is selected for prefill on SM120. The results show four possible backends: trtllm_ragged (TensorRT-LLM), fi (FlashInfer), cudnn (cuDNN), and fa (FlashAttention). The presence of _run_prefill_new_tokens_fa at line 2081, which calls _flash_attn_varlen_diff_headdims, indicates that FlashAttention is available as a prefill backend.

Assumptions Made

This message rests on several assumptions, most of which are well-justified but worth examining. The assistant assumes that the prefill path is the correct place to investigate — that garbage output from the model originates in the prompt processing phase rather than the decode phase. This is a reasonable assumption because the logprobs for early tokens (like the space after "1") are already anomalous, suggesting the hidden states are corrupted from the very first forward pass.

The assistant also assumes that the attention kernel selection logic in vLLM's MLA implementation is relevant to the bug. Specifically, the assistant assumes that different backends (FlashAttention, FlashInfer, TensorRT-LLM, cuDNN) may have different levels of support for SM120 Blackwell, and that one of them might be silently falling back to an incorrect implementation. This assumption is grounded in the earlier discovery that the Triton MLA backend's supported_kv_cache_dtypes only listed auto and bfloat16 — not float16 — which could indicate incomplete Blackwell support.

A more subtle assumption is that the kv_b_proj weight, despite the apparent shape mismatch between the full [28672, 512] tensor and the expected [3584, 512] sharded parameter, is actually being used in the prefill computation. The assistant implicitly trusts that if there were a shape mismatch error, it would have manifested as a runtime exception rather than silent garbage output. This assumption may be incorrect — as noted in the chunk summary, the parameter might be materialized as an UninitializedParameter that silently produces incorrect results.

Potential Mistakes and Incorrect Assumptions

The most significant potential mistake in this message is the framing itself. By asking "is the issue with flash attention for prefill on SM120?", the assistant may be prematurely narrowing the search space. The garbage output could equally well be caused by:

The kv_b_proj weight being incorrectly loaded or sharded (the original hypothesis)
The MoE expert weights being incorrectly unmerged or mapped
The indexer weights (which required force-dequantization) being loaded incorrectly
The layer normalization weights being wrong
The embedding or lm_head weights being corrupted The assistant has already eliminated some of these (the GGUF dequant kernel works, the name mapping is correct) but has not yet verified that the actual loaded weight values match the expected values. A more direct approach might be to dump the weights of a single layer and compare them against a reference implementation. Additionally, the assistant assumes that the prefill path is the relevant one, but the model could be generating garbage during the decode phase even if prefill works correctly. The logprob data from message [msg 1913] shows that even the first few tokens have very low logprobs, which does suggest prefill corruption, but this is not definitive proof.

Input Knowledge Required

To fully understand this message, one needs knowledge of several domains. First, the Multi-head Latent Attention (MLA) architecture used by DeepSeek-V2/V3 and GLM-5: in MLA, the key-value cache stores a compressed latent representation (kv_c) and a positional embedding (k_pe), rather than full K/V tensors. The kv_b_proj layer projects the compressed latent back to the full attention dimension. Understanding this architecture is essential to grasp why the kv_b_proj weight and its sharding are so critical.

Second, one needs knowledge of vLLM's attention backend architecture. vLLM supports multiple attention implementations (FlashAttention, FlashInfer, TensorRT-LLM, cuDNN, Triton) and selects one based on hardware capabilities and model configuration. The _run_prefill_new_tokens method is the dispatch point for prefill, and the specific backend chosen depends on what's available at runtime.

Third, knowledge of tensor parallelism is required. With 8 GPUs, each ColumnParallelLinear layer splits its weight matrix column-wise across GPUs. If the weight is loaded as a full matrix but the loader expects a sharded matrix, the results will be incorrect — but the mechanism of failure depends on how vLLM handles the mismatch.

Finally, understanding the Blackwell SM120 architecture and its implications for CUDA kernel compatibility is crucial. The assistant has been operating under the assumption that Blackwell may have incomplete support in various ML libraries, which motivates the investigation of the attention backend.

Output Knowledge Created

This message creates diagnostic knowledge about the prefill backend selection in vLLM's MLA implementation. The grep output reveals that the codebase supports four prefill backends (trtllm_ragged, fi, cudnn, fa) and that FlashAttention's _flash_attn_varlen_diff_headdims function is available. This information allows the assistant to proceed to the next logical step: checking which backend is actually selected at runtime on SM120, and whether that backend has known issues with Blackwell GPUs.

The message also implicitly documents the code structure of vLLM's MLA attention at a specific commit (v0.16.0rc2.dev313). The line numbers (1993, 1998, 2003, 2014, 2044, 2081, 2084, 2097) provide a map of the prefill dispatch logic that future debugging can reference.

Conclusion

Message [msg 1928] is a brief but consequential step in a complex debugging journey. It represents the moment when the investigation pivots from weight loading to computation correctness — from asking "are the weights in the right place?" to asking "is the computation using them correctly?" The assistant's reasoning is concise but logically complete, connecting the kv_b_proj weight concern to the attention kernel dispatch and formulating a testable hypothesis about FlashAttention on SM120. Whether the hypothesis proves correct or not, the systematic approach of enumerating components in the forward pass and checking each one is a model of disciplined debugging. The message also illustrates a broader truth about deploying large models on novel hardware: the most insidious bugs are not crashes but silent corruptions, where everything loads successfully yet the output is nonsense, and the only way forward is methodical isolation of each component in the computational graph.