biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix] Use null block (0) for padded block table entries (#35431)

Browse Source 
Signed-off-by: SandishKumarHN <sandish@fb.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
This commit is contained in:
SandishKumarHN
2026-03-30 14:02:51 -07:00
committed by GitHub
parent 1fc69f59bb
commit bcc6f67447
15 changed files with 125 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/model_executor/layers/lightning_attn.py
View File

@@ -5,6 +5,7 @@ import torch
from einops import rearrange
from vllm.triton_utils import tl, triton
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@triton.jit
@@ -602,6 +603,7 @@ def _linear_attn_decode_kernel(
cache_h_stride,
cache_d0_stride,
cache_d1_stride,
pad_slot_id: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""
@@ -616,8 +618,8 @@ def _linear_attn_decode_kernel(
# Load slot index for the current batch
slot_id = tl.load(slot_idx + pid_b).to(tl.int64)
# Skip if slot_id is -1 (padding)
if slot_id == -1:
# Skip if slot_id is PAD_SLOT_ID (padding)
if slot_id == pad_slot_id:
return
batch_id = pid_b
@@ -727,6 +729,7 @@ def linear_decode_forward_triton(
cache_h_stride,
cache_d0_stride,
cache_d1_stride,
pad_slot_id=PAD_SLOT_ID,
BLOCK_SIZE=BLOCK_SIZE,
)

39
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
View File

@@ -9,7 +9,7 @@ import numpy as np
import torch
from vllm.triton_utils import tl, triton
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
@triton.jit()
@@ -49,12 +49,13 @@ def _causal_conv1d_fwd_kernel( # continuous batching
stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M
# others
pad_slot_id: tl.constexpr,
null_block_id: tl.constexpr,
# Meta-parameters
HAS_BIAS: tl.constexpr,
KERNEL_WIDTH: tl.constexpr,
SILU_ACTIVATION: tl.constexpr,
IS_APC_ENABLED: tl.constexpr,
USE_PAD_SLOT: tl.constexpr,
HAS_NULL_BLOCK: tl.constexpr,
NP2_STATELEN: tl.constexpr,
BLOCK_M: tl.constexpr,
BLOCK_N: tl.constexpr,
@@ -133,9 +134,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching
conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
).to(tl.int64)
if USE_PAD_SLOT: # noqa
if conv_states_input_coord == pad_slot_id:
# not processing as this is not the actual sequence
if HAS_NULL_BLOCK: # noqa
if conv_states_input_coord == null_block_id:
# not processing as this is a null block (padding)
return
conv_states_base = (
conv_states_ptr
@@ -475,6 +476,7 @@ def causal_conv1d_fn(
has_initial_state: torch.Tensor | None = None,
activation: str | None = "silu",
pad_slot_id: int = PAD_SLOT_ID,
null_block_id: int = NULL_BLOCK_ID,
block_idx_first_scheduled_token: torch.Tensor | None = None,
block_idx_last_scheduled_token: torch.Tensor | None = None,
initial_state_idx: torch.Tensor | None = None,
@@ -730,12 +732,13 @@ def causal_conv1d_fn(
block_size_to_align // BLOCK_M,
# others
pad_slot_id,
null_block_id,
# META
HAS_BIAS=bias is not None,
KERNEL_WIDTH=width,
SILU_ACTIVATION=activation in ["silu", "swish"],
IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
USE_PAD_SLOT=pad_slot_id is not None,
HAS_NULL_BLOCK=null_block_id is not None,
NP2_STATELEN=np2_statelen,
# launch_cooperative_grid=True
BLOCK_M=BLOCK_M,
@@ -778,7 +781,7 @@ def _causal_conv1d_update_kernel(
stride_o_dim: tl.constexpr,
stride_o_token: tl.constexpr,
# others
pad_slot_id: tl.constexpr,
null_block_id: tl.constexpr,
# Meta-parameters
HAS_BIAS: tl.constexpr,
KERNEL_WIDTH: tl.constexpr,
@@ -787,7 +790,7 @@ def _causal_conv1d_update_kernel(
IS_APC_ENABLED: tl.constexpr,
IS_SPEC_DECODING: tl.constexpr,
NP2_STATELEN: tl.constexpr,
USE_PAD_SLOT: tl.constexpr,
HAS_NULL_BLOCK: tl.constexpr,
BLOCK_N: tl.constexpr,
):
# ruff: noqa: E501
@@ -811,8 +814,8 @@ def _causal_conv1d_update_kernel(
conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
).to(tl.int64)
if USE_PAD_SLOT: # noqa
if conv_states_input_coord == pad_slot_id:
if HAS_NULL_BLOCK: # noqa
if conv_states_input_coord == null_block_id:
# not processing as this is not the actual sequence
return
@@ -1076,7 +1079,7 @@ def causal_conv1d_update(
num_accepted_tokens: torch.Tensor | None = None,
query_start_loc: torch.Tensor | None = None,
max_query_len: int = -1,
pad_slot_id: int = PAD_SLOT_ID,
null_block_id: int = NULL_BLOCK_ID,
block_idx_last_scheduled_token: torch.Tensor | None = None,
initial_state_idx: torch.Tensor | None = None,
validate_data=False,
@@ -1111,16 +1114,16 @@ def causal_conv1d_update(
max_query_len: int
If query_start_loc is not None, this indicates the maximum query
length in the batch.
pad_slot_id: int
if conv_state_indices is passed, lets the kernel identify padded
entries that will not be processed,
for example: conv_state_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
null_block_id: int
Block ID used to identify padded entries in
conv_state_indices. Block 0 is the null block.
for example: conv_state_indices = [null_block_id, 1, 20, null_block_id]
in this case, the kernel will not process entries at
indices 0 and 3
out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x`
"""
if validate_data:
assert pad_slot_id is not None
assert null_block_id is not None
assert x.stride(1) == 1
if isinstance(activation, bool):
activation = "silu" if activation is True else None
@@ -1225,7 +1228,7 @@ def causal_conv1d_update(
stride_o_dim,
stride_o_token,
# others
pad_slot_id,
null_block_id,
# META
HAS_BIAS=bias is not None,
KERNEL_WIDTH=width,
@@ -1234,7 +1237,7 @@ def causal_conv1d_update(
IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
IS_SPEC_DECODING=num_accepted_tokens is not None,
NP2_STATELEN=np2_statelen,
USE_PAD_SLOT=pad_slot_id is not None,
HAS_NULL_BLOCK=null_block_id is not None,
BLOCK_N=256,
)
if unsqueeze:

32
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
View File

@@ -10,7 +10,7 @@ from packaging import version
from vllm import _custom_ops as ops
from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
from vllm.triton_utils import HAS_TRITON, tl, triton
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
@@ -75,7 +75,7 @@ def _selective_scan_update_kernel(
out_ptr,
state_batch_indices_ptr,
dst_state_batch_indices_ptr,
pad_slot_id,
null_block_id,
num_accepted_tokens_ptr,
cu_seqlens_ptr,
# Matrix dimensions
@@ -203,7 +203,7 @@ def _selective_scan_update_kernel(
mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)
if HAS_STATE_BATCH_INDICES:
mask &= state_batch_idx != pad_slot_id
mask &= state_batch_idx != null_block_id
state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32)
if HAS_DT_BIAS:
@@ -257,7 +257,7 @@ def _selective_scan_update_kernel(
if IS_SPEC_DECODING:
dst_idx_ptr = dst_state_batch_indices_ptr + i_t * stride_dst_state_indices_T
token_dst_idx = tl.load(dst_idx_ptr).to(tl.int64)
if token_dst_idx != pad_slot_id:
if token_dst_idx != null_block_id:
token_dst_ptrs = (
state_ptr_base
+ token_dst_idx * stride_state_batch
@@ -329,7 +329,7 @@ def selective_state_update(
dt_softplus=False,
state_batch_indices=None,
dst_state_batch_indices=None,
pad_slot_id=PAD_SLOT_ID,
null_block_id=NULL_BLOCK_ID,
out=None,
num_accepted_tokens=None,
cu_seqlens=None,
@@ -348,12 +348,12 @@ def selective_state_update(
D: (dim,) or (nheads, dim)
z: (batch, dim) or (batch, nheads, dim)
dt_bias: (dim,) or (nheads, dim)
pad_slot_id: int
if cache_indices is passed, lets the kernel identify padded
entries that will not be processed,
for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
in this case, the kernel will not process entries at
indices 0 and 3
null_block_id: int
if state_batch_indices is passed, lets the kernel identify
padded entries that will not be processed,
for example: state_batch_indices = [null_block_id, 1, 20,
null_block_id] in this case, the kernel will not process
entries at indices 0 and 3
out: Preallocated ssm output tensor. Assume same shape as x.
In-place updated.
num_accepted_tokens: (batch,)
@@ -488,7 +488,7 @@ def selective_state_update(
out,
state_batch_indices,
dst_state_batch_indices,
pad_slot_id,
null_block_id,
num_accepted_tokens,
cu_seqlens,
N,
@@ -550,7 +550,7 @@ def selective_scan_fn(
query_start_loc=None,
cache_indices=None,
has_initial_state=None,
pad_slot_id=PAD_SLOT_ID,
null_block_id=NULL_BLOCK_ID,
block_size=1024,
block_idx_first_scheduled_token=None,
block_idx_last_scheduled_token=None,
@@ -588,10 +588,10 @@ def selective_scan_fn(
indicate if the ssm_state at the corresponding index should be
used as initial state. Not providing argument assumes
there's no initial state
pad_slot_id: int
null_block_id: int
if cache_indices is passed, lets the kernel identify padding entries
that will not be processed,
for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
for example: cache_indices = [null_block_id, 1 ,20 ,null_block_id]
in this case, the kernel will not process entries at indices 0 and 3
block_size: int
The block size to align the cached states to
@@ -643,7 +643,7 @@ def selective_scan_fn(
cache_indices,
has_initial_state,
ssm_states,
pad_slot_id,
null_block_id,
block_size,
block_idx_first_scheduled_token,
block_idx_last_scheduled_token,