diff --git a/tests/ut/patch/worker/test_patch_mamba_utils_source.py b/tests/ut/patch/worker/test_patch_mamba_utils_source.py index c2a606c05c12..b0b0ee4bdcff 100644 --- a/tests/ut/patch/worker/test_patch_mamba_utils_source.py +++ b/tests/ut/patch/worker/test_patch_mamba_utils_source.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Source-level checks for the Ascend Mamba precision-kernel override.""" +"""Source-level checks for the Ascend Mamba precision-kernel overridess.""" from __future__ import annotations @@ -68,4 +68,4 @@ def test_patch_only_installs_existing_ascend_postprocess_kernel() -> None: assert "mamba_utils.postprocess_mamba_fused_kernel = postprocess_mamba_fused_kernel" in patch_source assert "MambaBase.bind_kv_cache" not in patch_source assert "mamba_utils._copy_mamba_state_block" not in patch_source - assert "mamba_utils.precopy_mamba_align_fused_kernel" not in patch_source + assert "mamba_utils.precopy_mamba_align_fused_kernel" in patch_source diff --git a/vllm_ascend/ops/triton/mamba/precopy.py b/vllm_ascend/ops/triton/mamba/precopy.py new file mode 100644 index 000000000000..2d47e869a5d9 --- /dev/null +++ b/vllm_ascend/ops/triton/mamba/precopy.py @@ -0,0 +1,226 @@ +# Adapted from vllm/v1/worker/mamba_utils.py. +# SPDX-License-Identifier: Apache-2.0 + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _copy_mamba_state_block( + state_idx, + bt_row_idx, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + COPY_BLOCK_SIZE: tl.constexpr, + CONV_STATE_DIM_FIRST: tl.constexpr, +): + """Copy one Mamba state block without casting pointers in copy loops. + + triton-ascend's AxisInfo analysis can abort on + ``(integer_address + loop_offset).to(pointer_type)``. Cast the base + addresses once and use pointer arithmetic inside the loops instead. + """ + state_base_addr = tl.load(state_base_addrs_ptr + state_idx) + state_block_stride = tl.load(state_block_strides_ptr + state_idx) + state_elem_size = tl.load(state_elem_sizes_ptr + state_idx) + state_inner_size = tl.load(state_inner_sizes_ptr + state_idx) + conv_width = tl.load(state_conv_widths_ptr + state_idx) + + group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64) + group_base_addr = tl.load(block_table_ptrs_ptr + group_idx) + block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32)) + block_table_base = block_table_typed + bt_row_idx * block_table_stride_req + + dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64) + dst_addr = state_base_addr + dest_block_id * state_block_stride + + is_conv_state = conv_width > 0 + + if CONV_STATE_DIM_FIRST and is_conv_state: + src_block_id = tl.load(block_table_base + src_col).to(tl.int64) + + dim_rows = tl.load(state_dim_row_count_ptr + state_idx) + row_stride = tl.load(state_dim_row_stride_ptr + state_idx) + + per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size + + bias_bytes = token_bias.to(tl.int64) * state_elem_size + + src_block_addr = state_base_addr + src_block_id * state_block_stride + + offsets = tl.arange(0, COPY_BLOCK_SIZE) + + for row in range(0, dim_rows): + row_src = (src_block_addr + row * row_stride + bias_bytes).to(tl.pointer_type(tl.uint8)) + + row_dst = (dst_addr + row * row_stride).to(tl.pointer_type(tl.uint8)) + + for offset in range(0, per_row_bytes, COPY_BLOCK_SIZE): + mask = offset + offsets < per_row_bytes + + data = tl.load( + row_src + offset + offsets, + mask=mask, + ) + + tl.store( + row_dst + offset + offsets, + data, + mask=mask, + ) + + return + + if is_conv_state: + src_block_id = tl.load(block_table_base + src_col).to(tl.int64) + + src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size + + src_addr = state_base_addr + src_block_id * state_block_stride + src_offset + + copy_size = (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size + + offsets = tl.arange(0, COPY_BLOCK_SIZE) + + src_ptr = src_addr.to(tl.pointer_type(tl.uint8)) + dst_ptr = dst_addr.to(tl.pointer_type(tl.uint8)) + + for offset in range(0, copy_size, COPY_BLOCK_SIZE): + mask = offset + offsets < copy_size + + data = tl.load( + src_ptr + offset + offsets, + mask=mask, + ) + + tl.store( + dst_ptr + offset + offsets, + data, + mask=mask, + ) + + return + + actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64) + + src_addr = state_base_addr + actual_src_block_id * state_block_stride + + copy_size = state_inner_size * state_elem_size + + copy_size_u64 = copy_size // 8 + + src_u64 = src_addr.to(tl.pointer_type(tl.uint64)) + dst_u64 = dst_addr.to(tl.pointer_type(tl.uint64)) + + offsets = tl.arange(0, COPY_BLOCK_SIZE) + + for offset in range(0, copy_size_u64, COPY_BLOCK_SIZE): + mask = offset + offsets < copy_size_u64 + + data = tl.load( + src_u64 + offset + offsets, + mask=mask, + ) + + tl.store( + dst_u64 + offset + offsets, + data, + mask=mask, + ) + + tail_start = copy_size_u64 * 8 + tail_bytes = copy_size - tail_start + + tail_offsets = tl.arange(0, 8) + + tail_src = (src_addr + tail_start).to(tl.pointer_type(tl.uint8)) + + tail_dst = (dst_addr + tail_start).to(tl.pointer_type(tl.uint8)) + + tail_mask = tail_offsets < tail_bytes + + tail_data = tl.load( + tail_src + tail_offsets, + mask=tail_mask, + ) + + tl.store( + tail_dst + tail_offsets, + tail_data, + mask=tail_mask, + ) + + +@triton.jit +def precopy_mamba_align_fused_kernel( + mamba_state_idx_ptr, + src_col_ptr, + token_bias_ptr, + block_table_ptrs_ptr, + block_table_stride_req: tl.int64, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + idx_mapping_ptr, + num_reqs, + COPY_BLOCK_SIZE: tl.constexpr, + CONV_STATE_DIM_FIRST: tl.constexpr, + HAS_IDX_MAPPING: tl.constexpr, +): + batch_idx = tl.program_id(0) + state_idx = tl.program_id(1) + + if batch_idx >= num_reqs: + return + + if HAS_IDX_MAPPING: + req_idx = tl.load(idx_mapping_ptr + batch_idx) + + if req_idx < 0: + return + else: + req_idx = batch_idx + + src_col = tl.load(src_col_ptr + req_idx) + + dst_col = tl.load(mamba_state_idx_ptr + req_idx) + + if src_col < 0 or src_col == dst_col: + return + + token_bias = tl.load(token_bias_ptr + req_idx) + + _copy_mamba_state_block( + state_idx, + batch_idx, + src_col, + dst_col, + token_bias, + block_table_ptrs_ptr, + block_table_stride_req, + state_base_addrs_ptr, + state_block_strides_ptr, + state_elem_sizes_ptr, + state_inner_sizes_ptr, + state_conv_widths_ptr, + state_group_indices_ptr, + state_dim_row_count_ptr, + state_dim_row_stride_ptr, + COPY_BLOCK_SIZE, + CONV_STATE_DIM_FIRST, + ) diff --git a/vllm_ascend/patch/worker/patch_mamba_utils.py b/vllm_ascend/patch/worker/patch_mamba_utils.py index ae359e070ccf..dbd5bfce9edf 100644 --- a/vllm_ascend/patch/worker/patch_mamba_utils.py +++ b/vllm_ascend/patch/worker/patch_mamba_utils.py @@ -16,6 +16,7 @@ from vllm_ascend.ops.triton.batch_memcpy import batch_memcpy_kernel from vllm_ascend.ops.triton.mamba.postprocess import postprocess_mamba_fused_kernel +from vllm_ascend.ops.triton.mamba.precopy import precopy_mamba_align_fused_kernel from vllm_ascend.utils import is_310p @@ -196,8 +197,10 @@ def _batch_memcpy_unavailable(src_ptrs, dst_ptrs, sizes): mamba_utils.batch_memcpy_kernel = batch_memcpy_kernel mamba_utils.batch_memcpy = _batch_memcpy_triton # Keep the existing Ascend postprocess precision fix. The shared copy - # helper and align pre-copy continue to use the upstream implementation. + # helper and align pre-copy use Ascend-safe implementations which hoist + # pointer casts out of copy loops. mamba_utils.postprocess_mamba_fused_kernel = postprocess_mamba_fused_kernel + mamba_utils.precopy_mamba_align_fused_kernel = precopy_mamba_align_fused_kernel else: mamba_utils.batch_memcpy = _batch_memcpy_unavailable mamba_utils.collect_mamba_copy_meta = _collect_mamba_copy_meta_torch diff --git a/vllm_ascend/worker/v2/attn_utils.py b/vllm_ascend/worker/v2/attn_utils.py index 5b72f2753090..f60f102b803f 100644 --- a/vllm_ascend/worker/v2/attn_utils.py +++ b/vllm_ascend/worker/v2/attn_utils.py @@ -227,6 +227,10 @@ def build_attn_metadata( if model_specific_attn_metadata is not None else {} ) + common_is_prefilling = common_attn_metadata_extra_kwargs.pop( + "is_prefilling", + is_prefilling, + ) common_attn_metadata = AscendCommonAttentionMetadata( query_start_loc=query_start_loc_gpu, query_start_loc_cpu=query_start_loc_cpu, @@ -242,7 +246,7 @@ def build_attn_metadata( attn_state=attn_state, graph_pad_size=graph_pad_size, num_input_tokens=num_input_tokens, - is_prefilling=is_prefilling, + is_prefilling=common_is_prefilling, max_seq_len=max_seq_len, causal=group_causal, **common_attn_metadata_extra_kwargs,