cg fix

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -413,10 +413,6 @@ def __init__(
             dtype=torch.int32,
             device=device,
         )
-        # Per-request cache_seqlens buffer (all set to topk_tokens)
-        self.decode_cache_seqlens_buffer = torch.full(
-            (max_num_seqs,), self.topk_tokens, dtype=torch.int32, device=device
-        )
         self.req_id_per_token_buffer = torch.empty(
             (vllm_config.scheduler_config.max_num_batched_tokens,),
             dtype=torch.int32,
@@ -548,9 +544,7 @@ def _build_fp8_extra_metadata(
             query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
             decode_query_len = (query_start_loc_cpu[1] - query_start_loc_cpu[0]).item()
 
-            # Per-request cache_seqlens: [topk_tokens] * num_decodes
-            decode_cache_seqlens = self.decode_cache_seqlens_buffer[:num_decodes]
-
+            decode_cache_seqlens = common_attn_metadata.seq_lens[:num_decodes]
             tile_scheduler_metadata, num_splits = get_mla_metadata(
                 cache_seqlens=decode_cache_seqlens,
                 num_q_tokens_per_head_k=decode_query_len * self.num_heads,

vllm/v1/attention/backends/utils.py

@@ -856,11 +856,15 @@ def split_decodes_and_prefills(
         return 0, num_reqs, 0, num_tokens
 
     if require_uniform:
+        # check if we are in a padded uniform batch; this is used for full-CGs, some
+        # requests may have a query lenght of 0 but since they are padding its fine
+        # to treat them as decodes (ensures num_decodes matches the captured size)
+        if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
+            assert num_reqs * query_lens[0] == num_tokens, "tokens not padded correctly"
+            return num_reqs, 0, num_tokens, 0  # all decodes
         is_prefill = query_lens != query_lens[0]
     else:
-        # 0-query len indicates a padded request; leave this at the back
-        # of the batch with the prefills
-        is_prefill = (query_lens > decode_threshold) | (query_lens == 0)
+        is_prefill = query_lens > decode_threshold
 
     if not torch.any(is_prefill):
         return num_reqs, 0, num_tokens, 0