enable varlen full cudagraph#3893
Conversation
meta-cla
Bot
added
the
CLA Signed
liangel-02
force-pushed
the
fix_varlen_cudagraph
branch
from
e967c82 to
4e61e94
Compare
| # (vLLM refreshes them before each replay) BEFORE the None check, else replay | ||
| # short-circuits to output.fill_(0) (zeroed attention). No-op outside capture. | ||
| attn_metadata, _, kv_cache, _ = get_attention_context(layer.layer_name) | ||
|
|
There was a problem hiding this comment.
could you help explain what's the difference between attn_metadata.num_actual_tokens and attn_metadata.max_query_len
There was a problem hiding this comment.
yeah attn_metadata.max_query_len is set to 1 at cuda graph capture time by vllm, so if you then replay this, your runtime max_query_len > 1 so you will get nans. but if you set it to essentially worst worst case and then run the fwd, this will be prevented. this would normally come at a perf cost but i copied vllm's flash attention linearized combine kernel in Dao-AILab/flash-attention#2692
There was a problem hiding this comment.
attn_metadata.max_query_len is set to 1 at cuda graph capture time by vllm
This sounds a bit suspicious -- naively reading max_query_len set to 1 means it's decoding only. Why would cuda graph capture (for prefill / decode / mixed prefill-decode) only run on decode-only tasks?
| class PyTorchVarlenAttentionMetadataBuilder(FlashAttentionMetadataBuilder): | ||
| _cudagraph_support = AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE | ||
| _cudagraph_support = AttentionCGSupport.ALWAYS |
There was a problem hiding this comment.
basically 1-line change https://github.com/pytorch/torchtitan/blob/main/torchtitan/experiments/rl/models/attention.py#L71 right?
| @@ -180,23 +154,18 @@ def forward( | |||
|
|
|||
| cu_seqlens_q = attn_metadata.query_start_loc | |||
| seqused_k = attn_metadata.seq_lens | |||
| max_seqlen_q = attn_metadata.max_query_len | |||
| # Pin max_seqlen_q to the total token count for safe cudagraph capture and use | |||
| # the linearize FA3 combine kernel to preserve performance. | |||
There was a problem hiding this comment.
| # the linearize FA3 combine kernel to preserve performance. | |
| # the linearized FA3 combine kernel to preserve performance. |
| # (vLLM refreshes them before each replay) BEFORE the None check, else replay | ||
| # short-circuits to output.fill_(0) (zeroed attention). No-op outside capture. | ||
| attn_metadata, _, kv_cache, _ = get_attention_context(layer.layer_name) | ||
|
|
There was a problem hiding this comment.
attn_metadata.max_query_len is set to 1 at cuda graph capture time by vllm
This sounds a bit suspicious -- naively reading max_query_len set to 1 means it's decoding only. Why would cuda graph capture (for prefill / decode / mixed prefill-decode) only run on decode-only tasks?
2 participants
No description provided.