diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 7262e5034..2b3a7d3c1 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -26,18 +26,6 @@ if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner from sgl_kernel import merge_state_v2 -from sgl_kernel.flash_attn import flash_attn_varlen_func as flash_attn_varlen_func_fa3 -from sgl_kernel.flash_attn import flash_attn_with_kvcache as flash_attn_with_kvcache_fa3 - -flash_attn_varlen_func = flash_attn_varlen_func_fa3 -flash_attn_with_kvcache = flash_attn_with_kvcache_fa3 - -from sglang.jit_kernel.flash_attention_v4 import ( - flash_attn_varlen_func as flash_attn_varlen_func_fa4, -) -from sglang.jit_kernel.flash_attention_v4 import ( - flash_attn_with_kvcache as flash_attn_with_kvcache_fa4, -) @dataclass @@ -89,223 +77,6 @@ class FlashAttentionMetadata: swa_spec_metadata: Optional[FlashAttentionMetadata] = None -# Copied from: -# https://github.com/houseroad/vllm/blob/4e45bfcaf928bdb9bd952b4ac922a3c205589ae8/vllm/v1/attention/backends/flash_attn.py -# -# Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into -# local attention blocks, where each block is passed to the attention kernel -# as an independent local ("virtual") batch item. -# -# For example, if are performing a chunked prefill a batch of 3 sequences: -# q_seqlens = [4, 10, 5] -# kv_seqlens = [6, 17, 9] -# Then normally for regular attention we would compute with an attention mask -# for batch idx 0 (q_seqlens = 4, kv_seqlens = 6) like: -# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6) -# k_toks > 0 1 2 3 4 5 -# q_toks v _____________ -# 0 | 1 1 1 -# 1 | 1 1 1 1 -# 2 | 1 1 1 1 1 -# 3 | 1 1 1 1 1 1 -# -# for local attention (with attn_chunk_size = 4) we would compute with an -# attention mask like: -# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6, attn_chunk_size = 4) -# k_toks > 0 1 2 3 4 5 -# q_toks v _____________ -# 0 | 1 1 1 -# 1 | 1 1 1 1 -# 2 | 1 -# 3 | 1 1 -# -# We can simulate this mask using standard flash-attention by breaking the -# sequences into local ("virtual") batches, where each local batch item is a -# local attention block, so in this case batch idx 0 would be broken up into: -# -# local-batch idx: 0 (q_seqlens = 2, kv_seqlens = 4) (batch 0) -# k_toks > 0 1 2 3 -# q_toks v _____________ -# 0 | 1 1 1 -# 1 | 1 1 1 1 -# local-batch idx: 1 (q_seqlens = 2, kv_seqlens = 2) (batch 0) -# k_toks > 4 5 -# q_toks v _____________ -# 2 | 1 -# 3 | 1 1 -# -# e.g. if we have: -# attn_chunk_size = 4 -# query_start_loc_np = [0, 4, 14, 19] (q_seqlens = [4, 10, 5]) -# Then this function would return: -# __b0__ ______b1______ __b2__ < orig batch indices -# q_seqlens_local = [ 2, 2, 1, 4, 4, 1, 4, 1] -# cu_seqlens_q_local = [0, 4, 6, 10, 14, 18, 19, 23, 24] -# seqlens_k_local = [ 4, 2, 4, 4, 4, 1, 4, 1] -# block_table_local : shape[local_virtual_batches, pages_per_local_batch] -def make_local_attention_virtual_batches( - attn_chunk_size: int, - query_start_loc_np: np.ndarray, - seq_lens_np: np.ndarray, - block_table: torch.Tensor, - page_size: int = 0, -) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]: - """ - Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into - local attention blocks, where each block is passed to the attention kernel - as an independent local ("virtual") batch item. - - Args: - attn_chunk_size: Size of local attention chunks - query_start_loc_np: Cumulative sum of query lengths (numpy array) - seq_lens_np: Sequence lengths (numpy array) - block_table: Block table for KV cache - page_size: Size of each page in the KV cache - - Returns: - seqlens_q_local: Query sequence lengths for local attention - cu_seqlens_q_local: Cumulative sum of query sequence lengths for local attention - seqlens_k_local: Key sequence lengths for local attention - block_table_local: Block table for local attention - """ - # Adjust attention_chunk_size based on the actual sequence length - # to avoid index out of bounds errors - max_seq_len = seq_lens_np.max() - effective_chunk_size = min(attn_chunk_size, max_seq_len) - # Make sure effective_chunk_size is divisible by page_size - effective_chunk_size = (effective_chunk_size // page_size) * page_size - if effective_chunk_size < page_size: - effective_chunk_size = page_size - attn_chunk_size = effective_chunk_size - - q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1] - actual_batch_size = seq_lens_np.shape[0] - - # Handle if we are starting in the middle of a local attention block, - # we assume q_seqlens > 0 (for all elements), for each batch idx we compute - # the number of tokens that are not in the first local attention block and - # then we can simply use a cdiv for the rest. - # For example if we have: - # attn_chunk_size = 4 - # q_seqlens = [4, 10, 5] - # k_seqlens = [6, 17, 9] - # Then we would get: - # new_tokens_in_first_block = [2, 1, 4] - # local_blocks = [2, 4, 2] - q_tokens_in_first_block = np.minimum( - attn_chunk_size - ((seq_lens_np - q_seqlens) % attn_chunk_size), q_seqlens - ).astype(np.int32) - tokens_in_last_block = attn_chunk_size + (seq_lens_np % -attn_chunk_size) - local_blocks = 1 + cdiv(q_seqlens - q_tokens_in_first_block, attn_chunk_size) - - # Once we know the number of local blocks we can compute the request spans - # for each batch idx, we can figure out the number of "virtual" requests we - # have to make, - # For the above example we would get: - # seqlens_q_local = [2, 2, 1, 4, 4, 1, 4, 1] - # - # First Get batched arange. (E.g., [2, 4, 2] -> [0, 1, 0, 1, 2, 3, 0, 1]) - # (TODO: max a utility to share this code with _prepare_inputs) - # arange step 1. [2, 4, 2] -> [2, 6, 8] - cu_num_blocks = np.cumsum(local_blocks) - virtual_batches = cu_num_blocks[-1] - # arange step 2. [2, 6, 8] -> [0, 0, 2, 2, 2, 2, 6, 6] - block_offsets = np.repeat(cu_num_blocks - local_blocks, local_blocks) - # arange step 3. [0, 1, 0, 1, 2, 3, 0, 1] - arange = np.arange(virtual_batches, dtype=np.int32) - block_offsets - # also compute reverse arange (i.e. [1, 0, 3, 2, 1, 0, 1, 0]) - rarange = np.repeat(local_blocks, local_blocks) - arange - 1 - # Then we can compute the seqlens_q_local, handling the fact that the - # first and last blocks could be partial - seqlens_q_local = np.repeat(q_seqlens - q_tokens_in_first_block, local_blocks) - # set the first block since this may be a partial block - seqlens_q_local[arange == 0] = q_tokens_in_first_block - # set the remaining blocks - seqlens_q_local[arange > 0] = np.minimum( - seqlens_q_local - attn_chunk_size * (arange - 1), attn_chunk_size - )[arange > 0] - - # convert from q_seqlens to cu_seqlens_q - cu_seqlens_q_local = np.pad(np.cumsum(seqlens_q_local), (1, 0)).astype(np.int32) - - # compute the seqlens_k_local, - # basically a full local attention block for all but the last block in each - # batch - # For our example this will be: - # seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1] - seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32) - seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block - - k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - ( - rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks) - ) - # For the example the local attention blocks start at: - # _b0_ _____b1_____ _b2_ - # k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8] - block_starts = k_seqstarts_absolute // page_size - - assert attn_chunk_size % page_size == 0, ( - f"attn_chunk_size {attn_chunk_size} is not " - f"divisible by page_size {page_size}" - ) - pages_per_local_batch = attn_chunk_size // page_size - - # Create a block_table for the local attention blocks - # For out example if we have a block-table like (assuming page_size=2): - # block_table = [ - # [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], < batch 0 - # [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], < batch 1 - # [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], < batch 2 - # ] - # Then for the local batches we would want a block-table like - # block_table_local = [ - # [ 0, 1 ], < local-batch 0, (batch 0, starting from k[0]) - # [ 2, 3 ], < local-batch 1, (batch 0, starting from k[4]) - # [ 12, 13 ], < local-batch 2, (batch 1, starting from k[4]) - # [ 14, 15 ], < local-batch 3, (batch 1, starting from k[8]) - # [ 16, 17 ], < local-batch 4, (batch 1, starting from k[12]) - # [ 18, 19 ], < local-batch 5, (batch 1, starting from k[16]) - # [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4]) - # [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8]) - # ] - block_indices = np.broadcast_to( - np.arange(pages_per_local_batch, dtype=np.int32), - (virtual_batches, pages_per_local_batch), - ) + np.expand_dims(block_starts, axis=1) - # Ensure block_indices doesn't exceed block_table dimensions - # This is a critical safety check that prevents index out of bounds errors - # when dealing with large sequences (>8192 tokens) or when the block_table - # dimensions are smaller than what would be needed for the full attention chunk size. - block_indices = block_indices.flatten().clip(max=block_table.shape[1] - 1) - batch_indices = np.repeat( - np.arange(actual_batch_size, dtype=np.int32), - local_blocks * pages_per_local_batch, - ) - - # NOTE: https://github.com/pytorch/pytorch/pull/160256 causes performance - # regression when using numpy arrays (batch and block indices) to index into - # torch tensor (block_table). As a workaround, convert numpy arrays to torch - # tensor first, which recovers perf. - batch_indices_torch = torch.from_numpy(batch_indices) - block_indices_torch = torch.from_numpy(block_indices) - block_table_local = block_table[batch_indices_torch, block_indices_torch].view( - virtual_batches, -1 - ) - - return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, block_table_local - - -def cdiv(a: int, b: int) -> int: - """Ceiling division.""" - return -(a // -b) - - -# TODO(hebiao064): remove this once we have a better way to handle the merge_state_v2 torch.compile issue -@torch._dynamo.disable() -def merge_state_v2_wrapper(o, s_a, o_exp, s_b): - return merge_state_v2(o, s_a, o_exp, s_b) - - class FlashAttentionBackend(AttentionBackend): """FlashAttention backend implementation. @@ -360,7 +131,6 @@ class FlashAttentionBackend(AttentionBackend): isinstance(model_runner.token_to_kv_pool, SWAKVPool) and model_runner.token_to_kv_pool.swa_layer_nums > 0 ) - if self.use_sliding_window_kv_pool: self.token_to_kv_pool = model_runner.token_to_kv_pool @@ -371,8 +141,6 @@ class FlashAttentionBackend(AttentionBackend): ) self.speculative_step_id = speculative_step_id - self.fa_impl_ver = fa_impl_ver - # Local attention settings self.has_local_attention = model_runner.model_config.is_local_attention_model if self.has_local_attention: @@ -388,6 +156,24 @@ class FlashAttentionBackend(AttentionBackend): self.sliding_window_size is not None and self.sliding_window_size > -1 ) + # Select version + self.fa_impl_ver = fa_impl_ver + if self.fa_impl_ver == 3: + from sgl_kernel.flash_attn import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + ) + elif self.fa_impl_ver == 4: + from sglang.jit_kernel.flash_attention_v4 import ( + flash_attn_varlen_func, + flash_attn_with_kvcache, + ) + else: + raise ValueError(f"Invalid version: {self.fa_impl_ver=}") + + self.flash_attn_varlen_func = flash_attn_varlen_func + self.flash_attn_with_kvcache = flash_attn_with_kvcache + # If num_splits == 0, we use a heuristic to automatically determine the number of splits. # We set nums splits to 1 if deterministic inference is enabled. # See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ for more details. @@ -830,19 +616,8 @@ class FlashAttentionBackend(AttentionBackend): and not is_swa_layer ) - flash_attn_varlen_func_base = flash_attn_varlen_func_fa3 - flash_attn_with_kvcache_base = flash_attn_with_kvcache_fa3 - - flash_attn_varlen_func = ( - flash_attn_varlen_func_fa4 - if self.fa_impl_ver == 4 - else flash_attn_varlen_func_base - ) - flash_attn_with_kvcache = ( - flash_attn_with_kvcache_fa4 - if self.fa_impl_ver == 4 - else flash_attn_with_kvcache_base - ) + flash_attn_varlen_func = self.flash_attn_varlen_func + flash_attn_with_kvcache = self.flash_attn_with_kvcache kwargs = {} if sinks is not None: @@ -1189,13 +964,7 @@ class FlashAttentionBackend(AttentionBackend): if sinks is not None: kwargs["sinks"] = sinks - flash_attn_with_kvcache_base = flash_attn_with_kvcache_fa3 - - flash_attn_with_kvcache = ( - flash_attn_with_kvcache_fa4 - if self.fa_impl_ver == 4 - else flash_attn_with_kvcache_base - ) + flash_attn_with_kvcache = self.flash_attn_with_kvcache k_descale, v_descale = None, None # only use kv scaling if: 1) fp8 kv is explicitly enabled, 2) RadixAttention @@ -3041,3 +2810,220 @@ def draft_decode_set_expand_metadata( positions = mask.cumsum(dim=1) - 1 num_seqs = cache_loc.shape[0] page_table[:num_seqs, :].scatter_(1, positions, cache_loc) + + +# Copied from: +# https://github.com/houseroad/vllm/blob/4e45bfcaf928bdb9bd952b4ac922a3c205589ae8/vllm/v1/attention/backends/flash_attn.py +# +# Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into +# local attention blocks, where each block is passed to the attention kernel +# as an independent local ("virtual") batch item. +# +# For example, if are performing a chunked prefill a batch of 3 sequences: +# q_seqlens = [4, 10, 5] +# kv_seqlens = [6, 17, 9] +# Then normally for regular attention we would compute with an attention mask +# for batch idx 0 (q_seqlens = 4, kv_seqlens = 6) like: +# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6) +# k_toks > 0 1 2 3 4 5 +# q_toks v _____________ +# 0 | 1 1 1 +# 1 | 1 1 1 1 +# 2 | 1 1 1 1 1 +# 3 | 1 1 1 1 1 1 +# +# for local attention (with attn_chunk_size = 4) we would compute with an +# attention mask like: +# batch idx: 0 (q_seqlens = 4, kv_seqlens = 6, attn_chunk_size = 4) +# k_toks > 0 1 2 3 4 5 +# q_toks v _____________ +# 0 | 1 1 1 +# 1 | 1 1 1 1 +# 2 | 1 +# 3 | 1 1 +# +# We can simulate this mask using standard flash-attention by breaking the +# sequences into local ("virtual") batches, where each local batch item is a +# local attention block, so in this case batch idx 0 would be broken up into: +# +# local-batch idx: 0 (q_seqlens = 2, kv_seqlens = 4) (batch 0) +# k_toks > 0 1 2 3 +# q_toks v _____________ +# 0 | 1 1 1 +# 1 | 1 1 1 1 +# local-batch idx: 1 (q_seqlens = 2, kv_seqlens = 2) (batch 0) +# k_toks > 4 5 +# q_toks v _____________ +# 2 | 1 +# 3 | 1 1 +# +# e.g. if we have: +# attn_chunk_size = 4 +# query_start_loc_np = [0, 4, 14, 19] (q_seqlens = [4, 10, 5]) +# Then this function would return: +# __b0__ ______b1______ __b2__ < orig batch indices +# q_seqlens_local = [ 2, 2, 1, 4, 4, 1, 4, 1] +# cu_seqlens_q_local = [0, 4, 6, 10, 14, 18, 19, 23, 24] +# seqlens_k_local = [ 4, 2, 4, 4, 4, 1, 4, 1] +# block_table_local : shape[local_virtual_batches, pages_per_local_batch] +def make_local_attention_virtual_batches( + attn_chunk_size: int, + query_start_loc_np: np.ndarray, + seq_lens_np: np.ndarray, + block_table: torch.Tensor, + page_size: int = 0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, torch.Tensor]: + """ + Take in `query_start_loc_np` and `seq_lens_np` and break the sequences into + local attention blocks, where each block is passed to the attention kernel + as an independent local ("virtual") batch item. + + Args: + attn_chunk_size: Size of local attention chunks + query_start_loc_np: Cumulative sum of query lengths (numpy array) + seq_lens_np: Sequence lengths (numpy array) + block_table: Block table for KV cache + page_size: Size of each page in the KV cache + + Returns: + seqlens_q_local: Query sequence lengths for local attention + cu_seqlens_q_local: Cumulative sum of query sequence lengths for local attention + seqlens_k_local: Key sequence lengths for local attention + block_table_local: Block table for local attention + """ + # Adjust attention_chunk_size based on the actual sequence length + # to avoid index out of bounds errors + max_seq_len = seq_lens_np.max() + effective_chunk_size = min(attn_chunk_size, max_seq_len) + # Make sure effective_chunk_size is divisible by page_size + effective_chunk_size = (effective_chunk_size // page_size) * page_size + if effective_chunk_size < page_size: + effective_chunk_size = page_size + attn_chunk_size = effective_chunk_size + + q_seqlens = query_start_loc_np[1:] - query_start_loc_np[:-1] + actual_batch_size = seq_lens_np.shape[0] + + # Handle if we are starting in the middle of a local attention block, + # we assume q_seqlens > 0 (for all elements), for each batch idx we compute + # the number of tokens that are not in the first local attention block and + # then we can simply use a cdiv for the rest. + # For example if we have: + # attn_chunk_size = 4 + # q_seqlens = [4, 10, 5] + # k_seqlens = [6, 17, 9] + # Then we would get: + # new_tokens_in_first_block = [2, 1, 4] + # local_blocks = [2, 4, 2] + q_tokens_in_first_block = np.minimum( + attn_chunk_size - ((seq_lens_np - q_seqlens) % attn_chunk_size), q_seqlens + ).astype(np.int32) + tokens_in_last_block = attn_chunk_size + (seq_lens_np % -attn_chunk_size) + local_blocks = 1 + cdiv(q_seqlens - q_tokens_in_first_block, attn_chunk_size) + + # Once we know the number of local blocks we can compute the request spans + # for each batch idx, we can figure out the number of "virtual" requests we + # have to make, + # For the above example we would get: + # seqlens_q_local = [2, 2, 1, 4, 4, 1, 4, 1] + # + # First Get batched arange. (E.g., [2, 4, 2] -> [0, 1, 0, 1, 2, 3, 0, 1]) + # (TODO: max a utility to share this code with _prepare_inputs) + # arange step 1. [2, 4, 2] -> [2, 6, 8] + cu_num_blocks = np.cumsum(local_blocks) + virtual_batches = cu_num_blocks[-1] + # arange step 2. [2, 6, 8] -> [0, 0, 2, 2, 2, 2, 6, 6] + block_offsets = np.repeat(cu_num_blocks - local_blocks, local_blocks) + # arange step 3. [0, 1, 0, 1, 2, 3, 0, 1] + arange = np.arange(virtual_batches, dtype=np.int32) - block_offsets + # also compute reverse arange (i.e. [1, 0, 3, 2, 1, 0, 1, 0]) + rarange = np.repeat(local_blocks, local_blocks) - arange - 1 + # Then we can compute the seqlens_q_local, handling the fact that the + # first and last blocks could be partial + seqlens_q_local = np.repeat(q_seqlens - q_tokens_in_first_block, local_blocks) + # set the first block since this may be a partial block + seqlens_q_local[arange == 0] = q_tokens_in_first_block + # set the remaining blocks + seqlens_q_local[arange > 0] = np.minimum( + seqlens_q_local - attn_chunk_size * (arange - 1), attn_chunk_size + )[arange > 0] + + # convert from q_seqlens to cu_seqlens_q + cu_seqlens_q_local = np.pad(np.cumsum(seqlens_q_local), (1, 0)).astype(np.int32) + + # compute the seqlens_k_local, + # basically a full local attention block for all but the last block in each + # batch + # For our example this will be: + # seqlens_k_local = [4, 2, 4, 4, 4, 1, 4, 1] + seqlens_k_local = np.full(cu_num_blocks[-1], attn_chunk_size, dtype=np.int32) + seqlens_k_local[cu_num_blocks - 1] = tokens_in_last_block + + k_seqstarts_absolute = np.repeat(seq_lens_np, local_blocks) - ( + rarange * attn_chunk_size + np.repeat(tokens_in_last_block, local_blocks) + ) + # For the example the local attention blocks start at: + # _b0_ _____b1_____ _b2_ + # k_seqstarts_absolute = [0, 4, 4, 8, 12, 16, 4, 8] + block_starts = k_seqstarts_absolute // page_size + + assert attn_chunk_size % page_size == 0, ( + f"attn_chunk_size {attn_chunk_size} is not " + f"divisible by page_size {page_size}" + ) + pages_per_local_batch = attn_chunk_size // page_size + + # Create a block_table for the local attention blocks + # For out example if we have a block-table like (assuming page_size=2): + # block_table = [ + # [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9], < batch 0 + # [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], < batch 1 + # [20, 21, 22, 23, 24, 25, 26, 27, 28, 29], < batch 2 + # ] + # Then for the local batches we would want a block-table like + # block_table_local = [ + # [ 0, 1 ], < local-batch 0, (batch 0, starting from k[0]) + # [ 2, 3 ], < local-batch 1, (batch 0, starting from k[4]) + # [ 12, 13 ], < local-batch 2, (batch 1, starting from k[4]) + # [ 14, 15 ], < local-batch 3, (batch 1, starting from k[8]) + # [ 16, 17 ], < local-batch 4, (batch 1, starting from k[12]) + # [ 18, 19 ], < local-batch 5, (batch 1, starting from k[16]) + # [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4]) + # [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8]) + # ] + block_indices = np.broadcast_to( + np.arange(pages_per_local_batch, dtype=np.int32), + (virtual_batches, pages_per_local_batch), + ) + np.expand_dims(block_starts, axis=1) + # Ensure block_indices doesn't exceed block_table dimensions + # This is a critical safety check that prevents index out of bounds errors + # when dealing with large sequences (>8192 tokens) or when the block_table + # dimensions are smaller than what would be needed for the full attention chunk size. + block_indices = block_indices.flatten().clip(max=block_table.shape[1] - 1) + batch_indices = np.repeat( + np.arange(actual_batch_size, dtype=np.int32), + local_blocks * pages_per_local_batch, + ) + + # NOTE: https://github.com/pytorch/pytorch/pull/160256 causes performance + # regression when using numpy arrays (batch and block indices) to index into + # torch tensor (block_table). As a workaround, convert numpy arrays to torch + # tensor first, which recovers perf. + batch_indices_torch = torch.from_numpy(batch_indices) + block_indices_torch = torch.from_numpy(block_indices) + block_table_local = block_table[batch_indices_torch, block_indices_torch].view( + virtual_batches, -1 + ) + + return seqlens_q_local, cu_seqlens_q_local, seqlens_k_local, block_table_local + + +def cdiv(a: int, b: int) -> int: + """Ceiling division.""" + return -(a // -b) + + +# TODO(hebiao064): remove this once we have a better way to handle the merge_state_v2 torch.compile issue +@torch._dynamo.disable() +def merge_state_v2_wrapper(o, s_a, o_exp, s_b): + return merge_state_v2(o, s_a, o_exp, s_b) diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 4c3f9e2f2..3fd45aac0 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -41,12 +41,12 @@ if _is_cuda: try: from sgl_kernel.flash_attn import flash_attn_varlen_func - from sglang.jit_kernel.flash_attention_v4 import ( - flash_attn_varlen_func as flash_attn_varlen_func_fa4, - ) - def flash_attn_func(*args, ver: int = 3, **kwargs): if ver == 4: + from sglang.jit_kernel.flash_attention_v4 import ( + flash_attn_varlen_func as flash_attn_varlen_func_fa4, + ) + return flash_attn_varlen_func_fa4(*args, **kwargs) return flash_attn_varlen_func(*args, **kwargs) @@ -57,8 +57,6 @@ if _is_cuda: if _is_npu: import torch_npu -_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip - from sglang.srt.distributed import ( split_tensor_along_last_dim, tensor_model_parallel_all_gather, @@ -78,6 +76,8 @@ from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import add_prefix, get_bool_env_var +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip + ROTARY_EMBED_CLASSES = { "normal": apply_rotary_pos_emb, }