Skip to content

vllm.model_executor.warmup.fa4_cutedsl_config

FA4 MLA prefill CuTeDSL compile warmup config.

Classes:

Functions:

FA4MLAPrefillCompileRequest dataclass

One compile-only FA4 MLA prefill request.

Source code in vllm/model_executor/warmup/fa4_cutedsl_config.py 
58
59
60
61
62
63
64
65
66
67
@dataclass(frozen=True)
class FA4MLAPrefillCompileRequest:
    """One compile-only FA4 MLA prefill request."""

    key: Hashable
    compile_spec: FlashAttentionCuTeDSLCompileSpec

    # Compile this request.
    def compile(self) -> None:
        self.compile_spec.compile()

iter_fa4_mla_prefill_compile_requests(ctx)

Yield compile requests for this fixed MLA backend.

FA4 dedupes duplicate atomic kernel selections in its own JIT cache.

Source code in vllm/model_executor/warmup/fa4_cutedsl_config.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def iter_fa4_mla_prefill_compile_requests(
    ctx: FA4MLAPrefillCompileContext,
) -> Iterator[FA4MLAPrefillCompileRequest]:
    """Yield compile requests for this fixed MLA backend.

    FA4 dedupes duplicate atomic kernel selections in its own JIT cache.
    """
    seen: set[Hashable] = set()
    for compile_spec in iter_fa4_mla_prefill_compile_specs(ctx):
        key = compile_spec.request_key()
        if key in seen:
            continue
        seen.add(key)
        yield FA4MLAPrefillCompileRequest(
            key=key,
            compile_spec=compile_spec,
        )

iter_fa4_mla_prefill_compile_specs(ctx)

Yield compile-only FA4 MLA prefill requests for this fixed setup.

Source code in vllm/model_executor/warmup/fa4_cutedsl_config.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def iter_fa4_mla_prefill_compile_specs(
    ctx: FA4MLAPrefillCompileContext,
) -> Iterator[FlashAttentionCuTeDSLCompileSpec]:
    """Yield compile-only FA4 MLA prefill requests for this fixed setup."""

    arch_family = _fa4_architecture_family_from_compute_capability(
        *torch.cuda.get_device_capability()
    )
    if not _supports_fa4_mla_prefill(ctx, arch_family):
        return

    from vllm.v1.attention.backends.fa_utils import (
        FlashAttentionCuTeDSLCompileSpec,
    )

    batch_size = FA4_MLA_PREFILL_COMPILE_BATCH_SIZE
    v_stride = None
    if not ctx.requires_v_padding:
        v_stride = (
            ctx.num_heads * ctx.kv_nope_head_dim,
            ctx.kv_nope_head_dim,
            1,
        )

    for _, max_seqlen_q, max_seqlen_k in _shape_probes_for_context(ctx, arch_family):
        total_q_tokens = batch_size * max_seqlen_q
        total_kv_tokens = batch_size * max_seqlen_k
        for causal in FA4_MLA_PREFILL_CAUSAL_OPTIONS:
            for return_lse in FA4_MLA_PREFILL_LSE_OPTIONS:
                yield FlashAttentionCuTeDSLCompileSpec(
                    q_shape=(total_q_tokens, ctx.num_heads, ctx.qk_head_dim),
                    k_shape=(total_kv_tokens, ctx.num_heads, ctx.qk_head_dim),
                    v_shape=(
                        total_kv_tokens,
                        ctx.num_heads,
                        ctx.effective_v_head_dim,
                    ),
                    v_stride=v_stride,
                    q_dtype=ctx.dtype,
                    cu_seqlens_q_shape=(batch_size + 1,),
                    cu_seqlens_k_shape=(batch_size + 1,),
                    max_seqlen_q=max_seqlen_q,
                    max_seqlen_k=max_seqlen_k,
                    softmax_scale=ctx.scale,
                    causal=causal,
                    return_softmax_lse=return_lse,
                    num_splits=ctx.num_splits,
                    fa_version=ctx.fa_version,
                )