[BugFix][Mamba] Fix staged-write kernel and duplicate is_prefilling argument

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

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,
+    )

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

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,