Skip to content

feat: support cross-attn cuda graph for whisper #21161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mickqian wants to merge 4 commits into from
Closed

feat: support cross-attn cuda graph for whisper #21161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c8375a3
feat: support cross-attn cuda graph for whisper
mickqian Mar 23, 2026
6ef3439
upd
mickqian Mar 23, 2026
54dc327
upd
mickqian Mar 23, 2026
c06d1dd
upd
mickqian Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion python/sglang/srt/layers/attention/flashinfer_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1053,11 +1053,17 @@ def update_cross_attention(
# Normal attention
paged_kernel_lens = seq_lens
kv_start_idx = encoder_lens
paged_kernel_lens_cpu = seq_lens_cpu
else:
# Cross attention
paged_kernel_lens = encoder_lens
kv_start_idx = torch.zeros_like(encoder_lens)
seq_lens_sum = encoder_lens.sum().item()
paged_kernel_lens_cpu = (

@mickqian mickqian Mar 23, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be the encoder length in encoder-decoder arch

encoder_lens.to(device="cpu", non_blocking=False)
if encoder_lens is not None
else None
)
Comment on lines +1062 to +1066

@gemini-code-assist gemini-code-assist Bot Mar 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The condition if encoder_lens is not None is checked twice, which is redundant. Consider simplifying this by assigning a default value to paged_kernel_lens_cpu outside the if-else block.

paged_kernel_lens_cpu = None
if wrapper_id == 0:
    # Normal attention
    paged_kernel_lens = seq_lens
    kv_start_idx = encoder_lens
    paged_kernel_lens_cpu = seq_lens_cpu
else:
    # Cross attention
    paged_kernel_lens = encoder_lens
    kv_start_idx = torch.zeros_like(encoder_lens)
    seq_lens_sum = encoder_lens.sum().item()
    paged_kernel_lens_cpu = encoder_lens.to(device="cpu", non_blocking=False) if encoder_lens is not None else None


self.call_begin_forward(
decode_wrappers[wrapper_id],
Expand All @@ -1067,7 +1073,7 @@ def update_cross_attention(
self.kv_indptr[wrapper_id],
kv_start_idx,
spec_info,
seq_lens_cpu=seq_lens_cpu,
seq_lens_cpu=paged_kernel_lens_cpu,
)

def call_begin_forward(
Expand Down
14 changes: 10 additions & 4 deletions python/sglang/srt/model_executor/cuda_graph_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,6 @@ def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:

@dataclass
class DecodeInputBuffers(ForwardInputBuffers):

input_ids: torch.Tensor
input_embeds: torch.Tensor
req_pool_indices: torch.Tensor
Expand Down Expand Up @@ -436,7 +435,6 @@ def patch_model(


def set_torch_compile_config():
import torch._dynamo.config
import torch._inductor.config

torch._inductor.config.coordinate_descent_tuning = True
Expand Down Expand Up @@ -589,8 +587,16 @@ def __init__(self, model_runner: ModelRunner):
if self.dllm_config is None
else self.dllm_config.block_size
)

self.encoder_len_fill_value = 0
if self.is_encoder_decoder:
self.encoder_len_fill_value = int(
getattr(
self.model_runner.model_config.hf_config,
"max_source_positions",
1,
)
)
Comment on lines +590 to +597

@gemini-code-assist gemini-code-assist Bot Mar 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Consider extracting this logic into a separate function to improve readability and maintainability. The function could be named _get_max_source_positions or similar.

    def _get_max_source_positions(self):
        return int(
            getattr(
                self.model_runner.model_config.hf_config,
                "max_source_positions",
                1,
            )
        )

    def __init__(self, model_runner: ModelRunner):
        ...
        self.encoder_len_fill_value = 0
        if self.is_encoder_decoder:
            self.encoder_len_fill_value = self._get_max_source_positions()

else:
self.encoder_len_fill_value = 0

if self.enable_torch_compile:
set_torch_compile_config()
Expand Down
200 changes: 92 additions & 108 deletions python/sglang/srt/models/whisper.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,65 +99,16 @@ def forward(
kv, _ = self.kv_proj(cross_hidden_states)
k, v = kv.split([self.kv_size, self.kv_size], dim=-1)
else:
k = torch.zeros_like(q)
v = torch.zeros_like(q)
k = None
v = None

q = q * self.scaling
num_heads = self.attn.tp_q_head_num
head_dim = self.attn.head_dim

q = q.view(-1, num_heads, head_dim)
k = k.view(-1, num_heads, head_dim)
v = v.view(-1, num_heads, head_dim)

q_len = q.shape[0]
kv_len = k.shape[0]

q = q.transpose(0, 1)
k = k.transpose(0, 1)
v = v.transpose(0, 1)

attn_weights = torch.bmm(q, k.transpose(1, 2))

# Apply block-diagonal mask for batched cross-attention
batch_size = forward_batch.batch_size if forward_batch else 1
if batch_size > 1 and kv_len > 0:
encoder_len_per_request = kv_len // batch_size
if encoder_len_per_request * batch_size == kv_len:
is_decode = forward_batch.forward_mode.is_decode()
if is_decode:
mask = torch.zeros(
(q_len, kv_len), device=q.device, dtype=torch.bool
)
for i in range(batch_size):
enc_start = i * encoder_len_per_request
enc_end = (i + 1) * encoder_len_per_request
mask[i, enc_start:enc_end] = True
attn_weights = attn_weights.masked_fill(
~mask.unsqueeze(0), float("-inf")
)
else:
seq_lens = forward_batch.seq_lens
if seq_lens is not None and len(seq_lens) == batch_size:
seq_lens_list = seq_lens.tolist()
mask = torch.zeros(
(q_len, kv_len), device=q.device, dtype=torch.bool
)
q_start = 0
for i, dec_len in enumerate(seq_lens_list):
enc_start = i * encoder_len_per_request
enc_end = (i + 1) * encoder_len_per_request
q_end = q_start + dec_len
mask[q_start:q_end, enc_start:enc_end] = True
q_start = q_end
attn_weights = attn_weights.masked_fill(
~mask.unsqueeze(0), float("-inf")
)

attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1)
attn_output = torch.bmm(attn_weights, v)
attn_output = attn_output.transpose(0, 1)
attn_output = attn_output.reshape(q_len, num_heads * head_dim)
q = q.view(-1, self.attn.tp_q_head_num, self.attn.head_dim)

@mickqian mickqian Mar 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

replaced hand-written cross-attn with native backend (RadixAttention)

if k is not None:
k = k.view(-1, self.attn.tp_k_head_num, self.attn.qk_head_dim)
v = v.view(-1, self.attn.tp_v_head_num, self.attn.v_head_dim)

attn_output = self.attn(q, k, v, forward_batch)
else:
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.chunk(chunks=3, dim=-1)
Expand Down Expand Up @@ -313,7 +264,6 @@ def forward(


class WhisperEncoder(torch.nn.Module):

def __init__(
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
):
Expand Down Expand Up @@ -361,7 +311,6 @@ def forward(


class WhisperDecoder(torch.nn.Module):

def __init__(
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
):
Expand Down Expand Up @@ -408,7 +357,6 @@ def forward(


class WhisperForConditionalGeneration(torch.nn.Module):

def __init__(
self, config: WhisperConfig, quant_config: Optional[QuantizationConfig] = None
):
Expand All @@ -420,7 +368,6 @@ def __init__(
)
self.logits_processor = LogitsProcessor(config)
self.config = config
self._encoder_cache = {}

def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
stacked_params_mapping = [
Expand Down Expand Up @@ -468,8 +415,74 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)

def pad_input_ids(self, input_ids: List[int], _mm_inputs: MultimodalInputs):
return input_ids
def _get_encoder_len(self, mm_inputs: Optional[MultimodalInputs]) -> int:
if mm_inputs is None or not mm_inputs.mm_items:
return 0
return self.config.max_source_positions

def _batch_audio_inputs(
self, forward_batch: ForwardBatch
) -> tuple[Optional[torch.Tensor], List[int]]:
if (
forward_batch.forward_mode.is_decode()
or not forward_batch.mm_inputs
or all(forward_batch.encoder_cached)
):
return None, []

batched_features = []
encoder_lens_need = []
for i, mm_input in enumerate(forward_batch.mm_inputs):
if (
forward_batch.encoder_cached[i]
or mm_input is None
or not mm_input.mm_items
):
continue

features = mm_input.mm_items[0].feature
if features.ndim == 2:
features = features.unsqueeze(0)

batched_features.append(features)
encoder_lens_need.append(int(forward_batch.encoder_lens[i].item()))

if not batched_features:
return None, []

return torch.cat(batched_features, dim=0), encoder_lens_need

def _flatten_encoder_outputs(
self, encoder_outputs: torch.Tensor, encoder_lens_need: List[int]
) -> torch.Tensor:
hidden_size = encoder_outputs.shape[-1]
total_encoder_len = sum(encoder_lens_need)
encoder_outputs_flat = torch.zeros(
total_encoder_len,
hidden_size,
device=encoder_outputs.device,
dtype=encoder_outputs.dtype,
)

i = start_pos = 0
for encoder_len in encoder_lens_need:
if encoder_len == 0:
continue
end_pos = start_pos + encoder_len
encoder_outputs_flat[start_pos:end_pos] = encoder_outputs[i, :encoder_len]
i += 1
start_pos += encoder_len

return encoder_outputs_flat

def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs):
encoder_len = self._get_encoder_len(mm_inputs)
if encoder_len == 0:
return input_ids

mm_inputs.num_image_tokens = encoder_len
pad_id = self.config.pad_token_id if self.config.pad_token_id is not None else 0
return [pad_id] * encoder_len + input_ids

def forward(
self,
Expand All @@ -479,52 +492,23 @@ def forward(
**kwargs: Any,
) -> LogitsProcessorOutput:
dtype = self.encoder.conv1.weight.dtype
is_decode = forward_batch.forward_mode.is_decode()

if is_decode:
encoder_outputs = None
if forward_batch.req_pool_indices is not None:
req_indices = forward_batch.req_pool_indices.tolist()
encoder_list = []
for req_idx in req_indices:
if req_idx in self._encoder_cache:
encoder_list.append(self._encoder_cache[req_idx])
if encoder_list:
encoder_outputs = torch.cat(encoder_list, dim=0)
else:
encoder_list = []
mm_inputs_list = forward_batch.mm_inputs if forward_batch.mm_inputs else []
req_indices = (
forward_batch.req_pool_indices.tolist()
if forward_batch.req_pool_indices is not None
else []
encoder_outputs = None
batched_audio_features, encoder_lens_need = self._batch_audio_inputs(
forward_batch
)
if batched_audio_features is not None:
encoder_position_ids = torch.arange(
self.config.max_source_positions,
device=batched_audio_features.device,
)
encoder_outputs = self.encoder(
batched_audio_features.to(dtype),
encoder_position_ids,
forward_batch,
)
encoder_outputs = self._flatten_encoder_outputs(
encoder_outputs, encoder_lens_need
)

for req_idx, mm_input in zip(req_indices, mm_inputs_list):
if mm_input is None or not mm_input.mm_items:
continue

features = mm_input.mm_items[0].feature
if features.ndim == 2:
features = features.unsqueeze(0)

encoder_len = features.shape[-1] // 2
encoder_position_ids = torch.arange(encoder_len).to(
features.device, non_blocking=True
)

req_encoder_outputs = self.encoder(
features.to(dtype), encoder_position_ids, forward_batch
)
req_encoder_outputs = req_encoder_outputs.squeeze(0)

self._encoder_cache[req_idx] = req_encoder_outputs
encoder_list.append(req_encoder_outputs)

if encoder_list:
encoder_outputs = torch.cat(encoder_list, dim=0)
else:
encoder_outputs = None

decoder_outputs = self.decoder(
input_ids, encoder_outputs, forward_batch, positions
Expand Down
26 changes: 22 additions & 4 deletions python/sglang/srt/server_args.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2256,12 +2256,30 @@ def _handle_attention_backend_compatibility(self):
self.speculative_algorithm is None
), "Speculative decoding is currently not supported with Flex Attention backend"

# Encoder-decoder models (e.g., Whisper)
# Whisper currently only supports cuda graph on the native flashinfer
# cross-attention path.
if model_config.is_encoder_decoder:
logger.warning(
"Cuda graph is disabled for encoder-decoder models (e.g., Whisper)"
model_architectures = model_config.hf_config.architectures or []
is_whisper = "WhisperForConditionalGeneration" in model_architectures
effective_prefill_backend = (
self.prefill_attention_backend or self.attention_backend
)
self.disable_cuda_graph = True
effective_decode_backend = (
self.decode_attention_backend or self.attention_backend
)
if not is_whisper:
logger.warning(
"Cuda graph is disabled for non-Whisper encoder-decoder models"
)
self.disable_cuda_graph = True
elif (
effective_prefill_backend != "flashinfer"
or effective_decode_backend != "flashinfer"
):
logger.warning(
"Cuda graph is disabled for Whisper unless both prefill and decode attention backends are flashinfer"
)
self.disable_cuda_graph = True
Comment on lines +2275 to +2282

@gemini-code-assist gemini-code-assist Bot Mar 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic to disable CUDA graph for Whisper models is duplicated here. Consider creating a function to encapsulate this logic and reuse it in both places to improve maintainability.

            def _disable_cuda_graph_for_whisper(self):
                model_architectures = model_config.hf_config.architectures or []
                is_whisper = "WhisperForConditionalGeneration" in model_architectures
                effective_prefill_backend = (
                    self.prefill_attention_backend or self.attention_backend
                )
                effective_decode_backend = (
                    self.decode_attention_backend or self.attention_backend
                )
                if not is_whisper:
                    logger.warning(
                        "Cuda graph is disabled for non-Whisper encoder-decoder models"
                    )
                    self.disable_cuda_graph = True
                elif (
                    effective_prefill_backend != "flashinfer"
                    or effective_decode_backend != "flashinfer"
                ):
                    logger.warning(
                        "Cuda graph is disabled for Whisper unless both prefill and decode attention backends are flashinfer"
                    )
                    self.disable_cuda_graph = True

        # Whisper currently only supports cuda graph on the native flashinfer
        # cross-attention path.
        if model_config.is_encoder_decoder:
            self._disable_cuda_graph_for_whisper()


# Major NVIDIA platforms backends
if (
Expand Down
Loading