[Perf][Hybrid] Vectorize _copy_mamba_state_block to uint64 for temporal (#48110)
Francesco Fusco committed
Jul 18, 2026 at 06:43 UTC
c71a583aa9f81400528e67e3d818f66b804e8340
1 file changed
+54
-16
vllm/v1/worker/mamba_utils.py
+54
-16
@@ -109,24 +109,45 @@ def _copy_mamba_state_block(
109
src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
110
num_elems_to_copy = (conv_width - token_bias).to(tl.int64) * state_inner_size
111
copy_size = num_elems_to_copy * state_elem_size
112
- else:
113
- # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]]
114
- actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(
115
- tl.int64
116
- )
117
- src_addr = state_base_addr + actual_src_block_id * state_block_stride
118
- # Use natural block data size (inner_size * elem_size), NOT
119
- # state_block_stride which is the page stride and can exceed the
120
- # actual data when the state tensor uses as_strided page padding.
121
- copy_size = state_inner_size * state_elem_size
112
+ offsets = tl.arange(0, COPY_BLOCK_SIZE)
113
+ for i in range(0, copy_size, COPY_BLOCK_SIZE):
114
+ mask = (i + offsets) < copy_size
115
+ curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8))
116
+ curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8))
117
+ data = tl.load(curr_src, mask=mask)
118
+ tl.store(curr_dst, data, mask=mask)
119
+ return
120
121
+ # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]]
122
+ actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(tl.int64)
123
+ src_addr = state_base_addr + actual_src_block_id * state_block_stride
124
+ # Use natural block data size (inner_size * elem_size), NOT
125
+ # state_block_stride which is the page stride and can exceed the
126
+ # actual data when the state tensor uses as_strided page padding.
127
+ copy_size = state_inner_size * state_elem_size
128
+
129
+ # Vectorize via uint64 (8B per thread → LDG.64/STG.64): both temporal
130
+ # and SD conv produce src/dst addresses aligned to a full token slice
131
+ # (inner_size * elem_size) and a copy_size that's a multiple of it,
132
+ # which is 8B-aligned for all state dtypes in use. A masked byte tail
133
+ # covers any remaining 0-7 bytes (only reachable for sub-8B slices).
134
+ copy_size_u64 = copy_size // 8
135
+ src_u64 = src_addr.to(tl.pointer_type(tl.uint64))
136
+ dst_u64 = dst_addr.to(tl.pointer_type(tl.uint64))
137
offsets = tl.arange(0, COPY_BLOCK_SIZE)
124
- for i in range(0, copy_size, COPY_BLOCK_SIZE):
125
- mask = (i + offsets) < copy_size
126
- curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8))
127
- curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8))
128
- data = tl.load(curr_src, mask=mask)
129
- tl.store(curr_dst, data, mask=mask)
138
+ for i in range(0, copy_size_u64, COPY_BLOCK_SIZE):
139
+ mask = (i + offsets) < copy_size_u64
140
+ data = tl.load(src_u64 + i + offsets, mask=mask)
141
+ tl.store(dst_u64 + i + offsets, data, mask=mask)
142
+
143
+ tail_start = copy_size_u64 * 8
144
+ tail_bytes = copy_size - tail_start
145
+ tail_off = tl.arange(0, 8)
146
+ tail_src = (src_addr + tail_start).to(tl.pointer_type(tl.uint8))
147
+ tail_dst = (dst_addr + tail_start).to(tl.pointer_type(tl.uint8))
148
+ tail_mask = tail_off < tail_bytes
149
+ tail_data = tl.load(tail_src + tail_off, mask=tail_mask)
150
+ tl.store(tail_dst + tail_off, tail_data, mask=tail_mask)
151
152
153
@triton.jit
@@ -674,6 +695,23 @@ class MambaSpecDecodeGPUContext:
695
self.state_inner_sizes[idx] = (
696
state[0].numel() if state.dim() > 1 else 1
697
)
698
+ # Temporal copies are vectorized with uint64
699
+ # loads/stores; base pointer and block stride must
700
+ # be 8B-aligned (tail loop handles copy_size % 8).
701
+ base_addr = state.data_ptr()
702
+ block_stride_bytes = block_stride_elems * state.element_size()
703
+ assert base_addr % 8 == 0, (
704
+ f"layer {layer_name}: state.data_ptr() = "
705
+ f"{base_addr:#x} is not 8B-aligned; "
706
+ f"_copy_mamba_state_block uint64 "
707
+ f"vectorization requires it"
708
+ )
709
+ assert block_stride_bytes % 8 == 0, (
710
+ f"layer {layer_name}: block stride = "
711
+ f"{block_stride_bytes}B is not 8B-aligned; "
712
+ f"_copy_mamba_state_block uint64 "
713
+ f"vectorization requires it"
714
+ )
715
716
self.state_group_indices[idx] = group_local_idx
717
idx += 1