A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 14:45 UTC
full cudagraph for flex-attn (#36298)
Signed-off-by: shunting314 <shunting@meta.com>
shunting314 committed
Apr 2, 2026 at 21:15 UTC
8b141ed8c328ef42ff9deeb3d6e26ea63b9550c1
4 files changed
+145
-11
tests/compile/fullgraph/test_full_cudagraph.py
-11
@@ -170,14 +170,3 @@ class TestFullCUDAGraph:
170
piecewise_res.outputs[0].text.lower()
171
== full_res.outputs[0].text.lower()
172
)
173
-
174
-
175
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda")
176
-def test_full_cudagraph_with_invalid_backend():
177
- # Flex_Attention is not supported with full cuda graph
178
- with pytest.raises(RuntimeError):
179
- LLM(
180
- model="Qwen/Qwen2-1.5B-Instruct",
181
- compilation_config=CompilationConfig(cudagraph_mode="FULL"),
182
- attention_config={"backend": "FLEX_ATTENTION"},
183
- )
tests/kernels/test_flex_attention.py
+53
@@ -26,6 +26,59 @@ MINIMUM_TORCH_VERSION = version.parse("2.7.0")
26
DIRECT_BUILD_VERSION = version.parse("2.9.dev0")
27
28
29
+@pytest.mark.skipif(
30
+ not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
31
+ reason="CUDA not available or PyTorch version < 2.7",
32
+)
33
+def test_flex_attention_full_cudagraphs(vllm_runner):
34
+ """Test the numerics for flex attention full cudagraphs support."""
35
+ model_name = "Qwen/Qwen2.5-1.5B-Instruct"
36
+ seed = 42
37
+ max_tokens = 24
38
+ num_logprobs = 5
39
+ prompts = [
40
+ "Hello, my name is",
41
+ "The president of the United States is",
42
+ "The capital of France is",
43
+ ]
44
+
45
+ # Run with flex attention eager
46
+ set_random_seed(seed)
47
+ with vllm_runner(
48
+ model_name,
49
+ runner="generate",
50
+ tensor_parallel_size=1,
51
+ num_gpu_blocks_override=128,
52
+ enforce_eager=True,
53
+ attention_config={"backend": "FLEX_ATTENTION"},
54
+ ) as llm_flex:
55
+ output_eager = llm_flex.generate_greedy_logprobs(
56
+ prompts, max_tokens, num_logprobs
57
+ )
58
+
59
+ # Run with flex attention compiled
60
+ set_random_seed(seed)
61
+ with vllm_runner(
62
+ model_name,
63
+ runner="generate",
64
+ tensor_parallel_size=1,
65
+ num_gpu_blocks_override=128,
66
+ enforce_eager=False,
67
+ gpu_memory_utilization=0.85,
68
+ attention_config={"backend": "FLEX_ATTENTION"},
69
+ ) as llm_default:
70
+ output_compile = llm_default.generate_greedy_logprobs(
71
+ prompts, max_tokens, num_logprobs
72
+ )
73
+
74
+ check_logprobs_close(
75
+ outputs_0_lst=output_eager,
76
+ outputs_1_lst=output_compile,
77
+ name_0="eager",
78
+ name_1="compile",
79
+ )
80
+
81
+
82
@pytest.mark.skipif(
83
not torch.cuda.is_available() or TORCH_VERSION < MINIMUM_TORCH_VERSION,
84
reason="CUDA not available or PyTorch version < 2.7",
vllm/v1/attention/backends/flex_attention.py
+91
@@ -30,6 +30,7 @@ from vllm.utils.math_utils import cdiv
30
from vllm.utils.torch_utils import is_quantized_kv_cache, is_torch_equal_or_newer
31
from vllm.v1.attention.backend import (
32
AttentionBackend,
33
+ AttentionCGSupport,
34
AttentionImpl,
35
AttentionMetadataBuilder,
36
AttentionType,
@@ -315,6 +316,18 @@ class BlockSparsityHint(NamedTuple):
316
hint_fn: _block_sparsity_hint_signature
317
318
319
+def copy_to_persistent(dst, src):
320
+ try:
321
+ dst = dst.as_strided(src.shape, src.stride())
322
+ except RuntimeError as e:
323
+ raise RuntimeError(
324
+ f"Fail to re-stride a persistent tensor of shape {dst.shape} "
325
+ f"for a tensor of shape {src.shape}"
326
+ ) from e
327
+ dst.copy_(src)
328
+ return dst
329
+
330
+
331
@dataclass
332
class FlexAttentionMetadata:
333
causal: bool
@@ -340,6 +353,9 @@ class FlexAttentionMetadata:
353
physical_to_logical: torch.Tensor
354
decode_offset: torch.Tensor
355
num_blocks_per_seq: torch.Tensor
356
+ persistent_kv_indices: torch.Tensor
357
+ persistent_kv_num_blocks: torch.Tensor
358
+ persistent_doc_ids: torch.Tensor
359
360
# For logging.
361
num_input_tokens: int = 0 # Number of tokens including padding.
@@ -656,8 +672,11 @@ class FlexAttentionMetadata:
672
kv_indices = unique_static_unsorted(
673
(used_pages_padded.long()), M=self.num_blocks
674
).to(torch.int32)
675
+ kv_indices = copy_to_persistent(self.persistent_kv_indices, kv_indices)
676
677
kv_num_blocks = (kv_indices >= 0).sum(dim=-1).to(torch.int32)
678
+ kv_num_blocks = copy_to_persistent(self.persistent_kv_num_blocks, kv_num_blocks)
679
+
680
block_mask_kwargs = {
681
"seq_lengths": (self.num_actual_tokens, self.total_cache_tokens),
682
"kv_num_blocks": kv_num_blocks[None, None],
@@ -694,6 +713,7 @@ class FlexAttentionMetadata:
713
assert self.suffix_kv_lens is None, "Not implemented yet."
714
# Create a lookup mapping from query indices -> request number
715
self.doc_ids = _offsets_to_doc_ids_tensor(self.query_start_loc)
716
+ self.doc_ids = copy_to_persistent(self.persistent_doc_ids, self.doc_ids)
717
self.num_blocks = self.total_cache_tokens // self.block_size
718
719
self.mask_mod = self.get_mask_mod()
@@ -701,6 +721,8 @@ class FlexAttentionMetadata:
721
722
723
class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadata]):
724
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
725
+
726
def __init__(
727
self,
728
kv_cache_spec: AttentionSpec,
@@ -726,6 +748,38 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
748
self.q_block_size: int = 16 if supports_small_blocks else 128
749
self.kv_block_size: int = self.block_size if supports_small_blocks else 128
750
751
+ self.max_model_len = self.model_config.max_model_len
752
+ max_num_seqs = vllm_config.scheduler_config.max_num_seqs
753
+ max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
754
+ self.max_num_q_block = (
755
+ self.max_model_len + self.q_block_size - 1
756
+ ) // self.q_block_size
757
+ self.persistent_kv_num_blocks = torch.empty(
758
+ self.max_num_q_block, dtype=torch.int32, device=device
759
+ )
760
+ self.persistent_offset_tensor = torch.empty(
761
+ max_num_seqs, dtype=torch.int32, device=device
762
+ )
763
+ self.persistent_doc_ids = torch.empty(
764
+ max_num_batched_tokens, dtype=torch.int32, device=device
765
+ )
766
+
767
+ # initialize later when we can access block_table
768
+ self.persistent_physical_to_logical = None
769
+ self.persistent_kv_indices = None
770
+
771
+ def build_for_cudagraph_capture(
772
+ self, common_attn_metadata: CommonAttentionMetadata
773
+ ) -> FlexAttentionMetadata:
774
+ # Use actual max_seq_len instead of max_model_len to avoid
775
+ # torch.compile recompilation during CUDA graph capture.
776
+ common_attn_metadata.max_seq_len = (
777
+ common_attn_metadata.seq_lens_cpu.max().item()
778
+ )
779
+ return self.build(
780
+ common_prefix_len=0, common_attn_metadata=common_attn_metadata
781
+ )
782
+
783
def build(
784
self,
785
common_prefix_len: int,
@@ -765,8 +819,32 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
819
inverse_block_table = physical_to_logical_mapping(
820
block_table_tensor, seq_lens, block_size, num_gpu_blocks
821
)
822
+ if self.persistent_physical_to_logical is None:
823
+ max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs
824
+ self.persistent_physical_to_logical = torch.empty(
825
+ max_num_seqs,
826
+ num_gpu_blocks,
827
+ dtype=torch.long,
828
+ device=self.device,
829
+ )
830
+
831
+ if self.persistent_kv_indices is None:
832
+ max_num_kv_block = (
833
+ self.max_model_len + self.kv_block_size - 1
834
+ ) // self.kv_block_size
835
+ self.persistent_kv_indices = torch.empty(
836
+ self.max_model_len,
837
+ max_num_kv_block,
838
+ dtype=torch.int32,
839
+ device=self.device,
840
+ )
841
+
842
+ inverse_block_table = copy_to_persistent(
843
+ self.persistent_physical_to_logical, inverse_block_table
844
+ )
845
846
offset_tensor = common_attn_metadata.compute_num_computed_tokens()
847
+ offset_tensor = copy_to_persistent(self.persistent_offset_tensor, offset_tensor)
848
849
out = FlexAttentionMetadata(
850
causal=common_attn_metadata.causal,
@@ -795,7 +873,20 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
873
direct_build=(self.direct_build and common_attn_metadata.causal),
874
q_block_size=self.q_block_size,
875
kv_block_size=self.kv_block_size,
876
+ persistent_kv_indices=self.persistent_kv_indices,
877
+ persistent_kv_num_blocks=self.persistent_kv_num_blocks,
878
+ persistent_doc_ids=self.persistent_doc_ids,
879
)
880
+
881
+ # Pre-build block_mask so it is ready before CUDA graph capture.
882
+ # Without this, the lazy build in forward() would run non-graph-safe
883
+ # ops (e.g. torch.nonzero) inside capture.
884
+ if out.block_mask is None:
885
+ if out.direct_build:
886
+ out.block_mask = out._build_block_mask_direct()
887
+ else:
888
+ out.block_mask = out.build_block_mask()
889
+
890
return out
891
892
def use_cascade_attention(self, *args, **kwargs) -> bool:
vllm/v1/worker/gpu_model_runner.py
+1
@@ -6077,6 +6077,7 @@ class GPUModelRunner(
6077
skip_eplb=True,
6078
remove_lora=False,
6079
num_active_loras=desc.num_active_loras,
6080
+ profile_seq_lens=profile_seq_lens,
6081
)
6082
self._dummy_run(
6083
desc.num_tokens,