@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / bcc6f67447
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 19:45 UTC

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

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>

SandishKumarHN committed Mar 30, 2026 at 14:02 UTC bcc6f67447e910c57b6986500cd7626e0974f8ef
15 files changed +125 -100
csrc/mamba/mamba_ssm/selective_scan.h
+1 -1
@@ -21,7 +21,7 @@ struct SSMParamsBase {
21 int dim_ngroups_ratio;
22 bool is_variable_B;
23 bool is_variable_C;
24 - int64_t pad_slot_id;
24 + int64_t null_block_id;
25
26 bool delta_softplus;
27 bool cache_enabled;
csrc/mamba/mamba_ssm/selective_scan_fwd.cu
+15 -7
@@ -118,9 +118,17 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
118
119 const int* cache_indices = params.cache_indices_ptr == nullptr ? nullptr
120 : reinterpret_cast<int *>(params.cache_indices_ptr);
121 - const int cache_index = cache_indices == nullptr ? batch_id : cache_indices[batch_id];
122 - // cache_index == params.pad_slot_id is defined as padding, so we exit early
123 - if (cache_index == params.pad_slot_id){
121 + int cache_index;
122 + if (cache_indices == nullptr) {
123 + cache_index = batch_id;
124 + } else if (params.cache_enabled) {
125 + const int* initial_state_idx = reinterpret_cast<const int*>(params.initial_state_idx_ptr);
126 + cache_index = cache_indices[batch_id * params.cache_indices_stride + initial_state_idx[batch_id]];
127 + } else {
128 + cache_index = cache_indices[batch_id];
129 + }
130 + // Skip batch entries whose cache index maps to the null block (padding).
131 + if (cache_indices != nullptr && cache_index == params.null_block_id){
132 return;
133 }
134 input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + sequence_start_index * params.u_batch_stride
@@ -527,7 +535,7 @@ void set_ssm_params_fwd(SSMParamsBase &params,
535 const std::optional<at::Tensor>& cache_indices,
536 const std::optional<at::Tensor>& has_initial_state,
537 bool varlen,
530 - int64_t pad_slot_id,
538 + int64_t null_block_id,
539 int64_t block_size,
540 const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
541 const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
@@ -544,7 +552,7 @@ void set_ssm_params_fwd(SSMParamsBase &params,
552 params.dstate = dstate;
553 params.n_groups = n_groups;
554 params.dim_ngroups_ratio = dim / n_groups;
547 - params.pad_slot_id = pad_slot_id;
555 + params.null_block_id = null_block_id;
556
557 params.delta_softplus = delta_softplus;
558
@@ -658,7 +666,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
666 const torch::Tensor &ssm_states,
667 // used to identify padding entries if cache_indices provided
668 // in case of padding, the kernel will return early
661 - int64_t pad_slot_id,
669 + int64_t null_block_id,
670 int64_t block_size,
671 const std::optional<torch::Tensor> &block_idx_first_scheduled_token,
672 const std::optional<torch::Tensor> &block_idx_last_scheduled_token,
@@ -805,7 +813,7 @@ void selective_scan_fwd(const torch::Tensor &u, const torch::Tensor &delta,
813 cache_indices,
814 has_initial_state,
815 varlen,
808 - pad_slot_id,
816 + null_block_id,
817 block_size,
818 block_idx_first_scheduled_token,
819 block_idx_last_scheduled_token,
csrc/ops.h
+1 -1
@@ -298,7 +298,7 @@ void selective_scan_fwd(
298 const std::optional<torch::Tensor>& query_start_loc,
299 const std::optional<torch::Tensor>& cache_indices,
300 const std::optional<torch::Tensor>& has_initial_state,
301 - const torch::Tensor& ssm_states, int64_t pad_slot_id, int64_t block_size,
301 + const torch::Tensor& ssm_states, int64_t null_block_id, int64_t block_size,
302 const std::optional<torch::Tensor>& block_idx_first_scheduled_token,
303 const std::optional<torch::Tensor>& block_idx_last_scheduled_token,
304 const std::optional<torch::Tensor>& initial_state_idx,
csrc/torch_bindings.cpp
+1 -1
@@ -556,7 +556,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
556 "Tensor? cache_indices,"
557 "Tensor? has_initial_state,"
558 "Tensor! ssm_states,"
559 - "int pad_slot_id,"
559 + "int null_block_id,"
560 "int block_size,"
561 "Tensor? block_idx_first_scheduled_token,"
562 "Tensor? block_idx_last_scheduled_token,"
tests/kernels/mamba/test_causal_conv1d.py
+24 -16
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
12 causal_conv1d_update,
13 )
14 from vllm.utils.torch_utils import set_random_seed
15 -from vllm.v1.attention.backends.utils import PAD_SLOT_ID
15 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
16
17
18 def causal_conv1d_ref(
@@ -122,7 +122,7 @@ def causal_conv1d_opcheck_fn(
122 has_initial_state: torch.Tensor | None = None,
123 conv_states: torch.Tensor | None = None,
124 activation: str | None = "silu",
125 - pad_slot_id: int = PAD_SLOT_ID,
125 + null_block_id: int = NULL_BLOCK_ID,
126 ):
127 """
128 x: (batch, dim, seqlen)
@@ -158,15 +158,16 @@ def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, ity
158 batch = 2
159 x = torch.randn(batch, dim, seqlen, device=device, dtype=itype)
160 x_ref = x.clone()
161 - conv_state = torch.randn(batch, dim, width - 1, device=device, dtype=itype)
161 + # +1 entry to reserve index 0 as null block
162 + conv_state = torch.randn(batch + 1, dim, width - 1, device=device, dtype=itype)
163
164 weight = torch.randn(dim, width, device=device, dtype=itype)
165 bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
165 - conv_state_ref = conv_state.detach().clone()
166 + # Start indices from 1, skipping null block at index 0
167 + conv_state_indices = torch.arange(1, batch + 1, dtype=torch.int32, device=device)
168 + conv_state_ref = conv_state[conv_state_indices].detach().clone()
169 activation = None if not silu_activation else "silu"
170
168 - conv_state_indices = torch.arange(batch, dtype=torch.int32, device=device)
169 -
171 out = causal_conv1d_update(
172 x,
173 conv_state,
@@ -179,7 +180,7 @@ def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation, ity
180 x_ref, conv_state_ref, weight, bias, activation=activation
181 )
182
182 - assert torch.equal(conv_state, conv_state_ref)
183 + assert torch.equal(conv_state[conv_state_indices], conv_state_ref)
184 assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
185
186
@@ -215,7 +216,8 @@ def test_causal_conv1d_update_with_batch_gather(
216
217 x_ref = x.clone()
218
218 - conv_state_indices = torch.randperm(total_entries)[:batch_size].to(
219 + # +1 to exclude index 0 (null block)
220 + conv_state_indices = (torch.randperm(total_entries - 1)[:batch_size] + 1).to(
221 dtype=torch.int32, device=device
222 )
223 unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device)
@@ -223,7 +225,9 @@ def test_causal_conv1d_update_with_batch_gather(
225 padded_state_indices = torch.concat(
226 [
227 conv_state_indices,
226 - torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
228 + torch.as_tensor(
229 + [NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device
230 + ),
231 ],
232 dim=0,
233 )
@@ -248,7 +252,6 @@ def test_causal_conv1d_update_with_batch_gather(
252 bias,
253 activation=activation,
254 conv_state_indices=padded_state_indices,
251 - pad_slot_id=PAD_SLOT_ID,
255 )
256 out_ref = causal_conv1d_update_ref(
257 x_ref[:batch_size], conv_state_ref, weight, bias, activation=activation
@@ -317,13 +320,19 @@ def test_causal_conv1d_varlen(
320 has_initial_states = torch.randint(
321 0, 2, (cumsum.shape[0] - 1,), dtype=torch.bool, device=x.device
322 )
320 - state_indices = torch.randperm(total_entries, dtype=torch.int32, device=x.device)[
321 - :batch_size
322 - ]
323 + # +1 to exclude index 0 (null block)
324 + state_indices = (
325 + torch.randperm(total_entries - 1, dtype=torch.int32, device=x.device)[
326 + :batch_size
327 + ]
328 + + 1
329 + )
330 padded_state_indices = torch.concat(
331 [
332 state_indices,
326 - torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
333 + torch.as_tensor(
334 + [NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device
335 + ),
336 ],
337 dim=-1,
338 )
@@ -336,7 +345,6 @@ def test_causal_conv1d_varlen(
345 cache_indices=padded_state_indices,
346 has_initial_state=has_initial_states,
347 activation=activation,
339 - pad_slot_id=PAD_SLOT_ID,
348 )
349
350 out_ref = []
@@ -345,7 +353,7 @@ def test_causal_conv1d_varlen(
353 splits = [torch.split(var, seqlens[0], dim=-1) for var in (x_ref)]
354 for i in range(len(seqlens[0])):
355 x_s = [v[i].unsqueeze(0) for v in splits][0]
348 - if padded_state_indices[i] == PAD_SLOT_ID:
356 + if padded_state_indices[i] == NULL_BLOCK_ID:
357 continue
358 out_ref_b.append(
359 causal_conv1d_ref(
tests/kernels/mamba/test_mamba_ssm.py
+29 -22
@@ -14,7 +14,7 @@ from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
14 )
15 from vllm.platforms import current_platform
16 from vllm.utils.torch_utils import set_random_seed
17 -from vllm.v1.attention.backends.utils import PAD_SLOT_ID
17 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
18
19
20 def selective_state_update_ref(
@@ -179,7 +179,7 @@ def selective_scan_opcheck_fn(
179 cache_indices=None,
180 has_initial_state=None,
181 ssm_states=None,
182 - pad_slot_id=PAD_SLOT_ID,
182 + null_block_id=NULL_BLOCK_ID,
183 block_size=2048,
184 block_idx_first_scheduled_token=None,
185 block_idx_last_scheduled_token=None,
@@ -229,7 +229,7 @@ def selective_scan_opcheck_fn(
229 cache_indices,
230 has_initial_state,
231 ssm_states,
232 - pad_slot_id,
232 + null_block_id,
233 block_size,
234 block_idx_first_scheduled_token,
235 block_idx_last_scheduled_token,
@@ -351,7 +351,6 @@ def test_selective_scan(
351 has_initial_state=torch.ones(batch_size, device=u.device, dtype=torch.bool)
352 if c > 0
353 else None,
354 - pad_slot_id=PAD_SLOT_ID,
354 block_size=2048,
355 block_idx_first_scheduled_token=None,
356 block_idx_last_scheduled_token=None,
@@ -652,15 +651,21 @@ def test_selective_scan_varlen(
651 prev_state_shape, device=u.device, dtype=itype, requires_grad=False
652 )
653 prev_state_ref = prev_state.clone()
655 - state_indices = torch.randperm(total_entries, dtype=torch.int32, device=u.device)[
656 - :batch_size
657 - ]
654 + # +1 to exclude index 0 (null block)
655 + state_indices = (
656 + torch.randperm(total_entries - 1, dtype=torch.int32, device=u.device)[
657 + :batch_size
658 + ]
659 + + 1
660 + )
661 unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device)
662 unused_states_bool[state_indices] = False
663 padded_state_indices = torch.concat(
664 [
665 state_indices,
663 - torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
666 + torch.as_tensor(
667 + [NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device
668 + ),
669 ],
670 dim=-1,
671 )
@@ -690,7 +695,7 @@ def test_selective_scan_varlen(
695 ]
696 for i in range(len(seqlens[0])):
697 u_s, delta_s, B_s, C_s, z_s = (v[i].unsqueeze(0) for v in splits)
693 - if padded_state_indices[i] == PAD_SLOT_ID:
698 + if padded_state_indices[i] == NULL_BLOCK_ID:
699 continue
700 out_ref_s, _ = selective_scan_ref(
701 u_s,
@@ -758,7 +763,8 @@ def test_selective_state_update_with_batch_indices(
763 padded_batch_size = batch_size + padding
764 total_entries = 10 * batch_size
765 state = torch.randn(total_entries, dim, dstate, dtype=itype, device=device)
761 - state_indices = torch.randperm(total_entries)[:batch_size].to(
766 + # +1 to exclude index 0 (null block)
767 + state_indices = (torch.randperm(total_entries - 1)[:batch_size] + 1).to(
768 dtype=torch.int32, device=device
769 )
770 unused_states_bool = torch.ones(total_entries, dtype=torch.bool, device=device)
@@ -766,7 +772,9 @@ def test_selective_state_update_with_batch_indices(
772 padded_state_indices = torch.concat(
773 [
774 state_indices,
769 - torch.as_tensor([PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
775 + torch.as_tensor(
776 + [NULL_BLOCK_ID] * padding, dtype=torch.int32, device=device
777 + ),
778 ],
779 dim=0,
780 )
@@ -793,7 +801,6 @@ def test_selective_state_update_with_batch_indices(
801 dt_bias=dt_bias,
802 dt_softplus=True,
803 state_batch_indices=padded_state_indices,
796 - pad_slot_id=PAD_SLOT_ID,
804 out=out,
805 )
806 out_ref = selective_state_update_ref(
@@ -849,7 +856,8 @@ def test_selective_state_update_with_heads_with_batch_indices(
856 state = torch.randn(
857 total_entries, nheads, headdim, dstate, dtype=itype, device=device
858 )
852 - state_indices = torch.randperm(total_entries)[:batch_size].to(
859 + # +1 to exclude index 0 (null block)
860 + state_indices = (torch.randperm(total_entries - 1)[:batch_size] + 1).to(
861 dtype=torch.int32, device=device
862 )
863
@@ -887,7 +895,6 @@ def test_selective_state_update_with_heads_with_batch_indices(
895 dt_bias=dt_bias,
896 dt_softplus=True,
897 state_batch_indices=state_indices,
890 - pad_slot_id=PAD_SLOT_ID,
898 out=out,
899 )
900 out_ref = selective_state_update_ref(
@@ -935,17 +942,18 @@ def test_selective_state_update_with_num_accepted_tokens(
942 state = torch.randn(total_state_slots, dim, dstate, dtype=itype, device=device)
943
944 state_batch_indices = torch.full(
938 - (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device
945 + (batch_size, max_seq_len), NULL_BLOCK_ID, dtype=torch.int32, device=device
946 )
947 + # Start from 1 to exclude null block at index 0
948 initial_state_slots = torch.randint(
941 - 0, 15, (batch_size,), device=device, dtype=torch.int32
949 + 1, 15, (batch_size,), device=device, dtype=torch.int32
950 )
951 for seq_idx in range(batch_size):
952 token_pos = max(num_accepted_tokens[seq_idx].item() - 1, 0)
953 state_batch_indices[seq_idx, token_pos] = initial_state_slots[seq_idx]
954
955 dst_state_batch_indices = torch.full(
948 - (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device
956 + (batch_size, max_seq_len), NULL_BLOCK_ID, dtype=torch.int32, device=device
957 )
958 slot_offset = 15
959 dst_slots_map = {}
@@ -1013,7 +1021,6 @@ def test_selective_state_update_with_num_accepted_tokens(
1021 state_batch_indices=state_batch_indices,
1022 dst_state_batch_indices=dst_state_batch_indices,
1023 num_accepted_tokens=num_accepted_tokens,
1016 - pad_slot_id=PAD_SLOT_ID,
1024 )
1025
1026 assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
@@ -1061,18 +1068,19 @@ def test_selective_state_update_varlen_with_num_accepted(
1068 state = torch.randn(total_state_slots, dim, dstate, dtype=itype, device=device)
1069
1070 state_batch_indices = torch.full(
1064 - (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device
1071 + (batch_size, max_seq_len), NULL_BLOCK_ID, dtype=torch.int32, device=device
1072 )
1073
1074 + # Start from 1 to exclude null block at index 0
1075 initial_state_slots = torch.randint(
1068 - 0, 15, (batch_size,), device=device, dtype=torch.int32
1076 + 1, 15, (batch_size,), device=device, dtype=torch.int32
1077 )
1078 for seq_idx in range(batch_size):
1079 token_pos = max(num_accepted_tokens[seq_idx].item() - 1, 0)
1080 state_batch_indices[seq_idx, token_pos] = initial_state_slots[seq_idx]
1081
1082 dst_state_batch_indices = torch.full(
1075 - (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device
1083 + (batch_size, max_seq_len), NULL_BLOCK_ID, dtype=torch.int32, device=device
1084 )
1085
1086 slot_offset = 15
@@ -1138,7 +1146,6 @@ def test_selective_state_update_varlen_with_num_accepted(
1146 state_batch_indices=state_batch_indices,
1147 dst_state_batch_indices=dst_state_batch_indices,
1148 num_accepted_tokens=num_accepted_tokens,
1141 - pad_slot_id=PAD_SLOT_ID,
1149 )
1150
1151 for seq_idx in range(batch_size):
vllm/_custom_ops.py
+2 -2
@@ -2062,7 +2062,7 @@ def selective_scan_fwd(
2062 cache_indices: torch.Tensor | None,
2063 has_initial_state: torch.Tensor | None,
2064 ssm_states: torch.Tensor,
2065 - pad_slot_id: int,
2065 + null_block_id: int,
2066 block_size: int = 1024,
2067 block_idx_first_scheduled_token: torch.Tensor | None = None,
2068 block_idx_last_scheduled_token: torch.Tensor | None = None,
@@ -2084,7 +2084,7 @@ def selective_scan_fwd(
2084 cache_indices,
2085 has_initial_state,
2086 ssm_states,
2087 - pad_slot_id,
2087 + null_block_id,
2088 block_size,
2089 block_idx_first_scheduled_token,
2090 block_idx_last_scheduled_token,
vllm/model_executor/layers/lightning_attn.py
+5 -2
@@ -5,6 +5,7 @@ import torch
5 from einops import rearrange
6
7 from vllm.triton_utils import tl, triton
8 +from vllm.v1.attention.backends.utils import PAD_SLOT_ID
9
10
11 @triton.jit
@@ -602,6 +603,7 @@ def _linear_attn_decode_kernel(
603 cache_h_stride,
604 cache_d0_stride,
605 cache_d1_stride,
606 + pad_slot_id: tl.constexpr,
607 BLOCK_SIZE: tl.constexpr,
608 ):
609 """
@@ -616,8 +618,8 @@ def _linear_attn_decode_kernel(
618 # Load slot index for the current batch
619 slot_id = tl.load(slot_idx + pid_b).to(tl.int64)
620
619 - # Skip if slot_id is -1 (padding)
620 - if slot_id == -1:
621 + # Skip if slot_id is PAD_SLOT_ID (padding)
622 + if slot_id == pad_slot_id:
623 return
624
625 batch_id = pid_b
@@ -727,6 +729,7 @@ def linear_decode_forward_triton(
729 cache_h_stride,
730 cache_d0_stride,
731 cache_d1_stride,
732 + pad_slot_id=PAD_SLOT_ID,
733 BLOCK_SIZE=BLOCK_SIZE,
734 )
735
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+21 -18
@@ -9,7 +9,7 @@ import numpy as np
9 import torch
10
11 from vllm.triton_utils import tl, triton
12 -from vllm.v1.attention.backends.utils import PAD_SLOT_ID
12 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
13
14
15 @triton.jit()
@@ -49,12 +49,13 @@ def _causal_conv1d_fwd_kernel( # continuous batching
49 stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M
50 # others
51 pad_slot_id: tl.constexpr,
52 + null_block_id: tl.constexpr,
53 # Meta-parameters
54 HAS_BIAS: tl.constexpr,
55 KERNEL_WIDTH: tl.constexpr,
56 SILU_ACTIVATION: tl.constexpr,
57 IS_APC_ENABLED: tl.constexpr,
57 - USE_PAD_SLOT: tl.constexpr,
58 + HAS_NULL_BLOCK: tl.constexpr,
59 NP2_STATELEN: tl.constexpr,
60 BLOCK_M: tl.constexpr,
61 BLOCK_N: tl.constexpr,
@@ -133,9 +134,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching
134 conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index
135 ).to(tl.int64)
136
136 - if USE_PAD_SLOT: # noqa
137 - if conv_states_input_coord == pad_slot_id:
138 - # not processing as this is not the actual sequence
137 + if HAS_NULL_BLOCK: # noqa
138 + if conv_states_input_coord == null_block_id:
139 + # not processing as this is a null block (padding)
140 return
141 conv_states_base = (
142 conv_states_ptr
@@ -475,6 +476,7 @@ def causal_conv1d_fn(
476 has_initial_state: torch.Tensor | None = None,
477 activation: str | None = "silu",
478 pad_slot_id: int = PAD_SLOT_ID,
479 + null_block_id: int = NULL_BLOCK_ID,
480 block_idx_first_scheduled_token: torch.Tensor | None = None,
481 block_idx_last_scheduled_token: torch.Tensor | None = None,
482 initial_state_idx: torch.Tensor | None = None,
@@ -730,12 +732,13 @@ def causal_conv1d_fn(
732 block_size_to_align // BLOCK_M,
733 # others
734 pad_slot_id,
735 + null_block_id,
736 # META
737 HAS_BIAS=bias is not None,
738 KERNEL_WIDTH=width,
739 SILU_ACTIVATION=activation in ["silu", "swish"],
740 IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
738 - USE_PAD_SLOT=pad_slot_id is not None,
741 + HAS_NULL_BLOCK=null_block_id is not None,
742 NP2_STATELEN=np2_statelen,
743 # launch_cooperative_grid=True
744 BLOCK_M=BLOCK_M,
@@ -778,7 +781,7 @@ def _causal_conv1d_update_kernel(
781 stride_o_dim: tl.constexpr,
782 stride_o_token: tl.constexpr,
783 # others
781 - pad_slot_id: tl.constexpr,
784 + null_block_id: tl.constexpr,
785 # Meta-parameters
786 HAS_BIAS: tl.constexpr,
787 KERNEL_WIDTH: tl.constexpr,
@@ -787,7 +790,7 @@ def _causal_conv1d_update_kernel(
790 IS_APC_ENABLED: tl.constexpr,
791 IS_SPEC_DECODING: tl.constexpr,
792 NP2_STATELEN: tl.constexpr,
790 - USE_PAD_SLOT: tl.constexpr,
793 + HAS_NULL_BLOCK: tl.constexpr,
794 BLOCK_N: tl.constexpr,
795 ):
796 # ruff: noqa: E501
@@ -811,8 +814,8 @@ def _causal_conv1d_update_kernel(
814 conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init
815 ).to(tl.int64)
816
814 - if USE_PAD_SLOT: # noqa
815 - if conv_states_input_coord == pad_slot_id:
817 + if HAS_NULL_BLOCK: # noqa
818 + if conv_states_input_coord == null_block_id:
819 # not processing as this is not the actual sequence
820 return
821
@@ -1076,7 +1079,7 @@ def causal_conv1d_update(
1079 num_accepted_tokens: torch.Tensor | None = None,
1080 query_start_loc: torch.Tensor | None = None,
1081 max_query_len: int = -1,
1079 - pad_slot_id: int = PAD_SLOT_ID,
1082 + null_block_id: int = NULL_BLOCK_ID,
1083 block_idx_last_scheduled_token: torch.Tensor | None = None,
1084 initial_state_idx: torch.Tensor | None = None,
1085 validate_data=False,
@@ -1111,16 +1114,16 @@ def causal_conv1d_update(
1114 max_query_len: int
1115 If query_start_loc is not None, this indicates the maximum query
1116 length in the batch.
1114 - pad_slot_id: int
1115 - if conv_state_indices is passed, lets the kernel identify padded
1116 - entries that will not be processed,
1117 - for example: conv_state_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
1117 + null_block_id: int
1118 + Block ID used to identify padded entries in
1119 + conv_state_indices. Block 0 is the null block.
1120 + for example: conv_state_indices = [null_block_id, 1, 20, null_block_id]
1121 in this case, the kernel will not process entries at
1122 indices 0 and 3
1123 out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x`
1124 """
1125 if validate_data:
1123 - assert pad_slot_id is not None
1126 + assert null_block_id is not None
1127 assert x.stride(1) == 1
1128 if isinstance(activation, bool):
1129 activation = "silu" if activation is True else None
@@ -1225,7 +1228,7 @@ def causal_conv1d_update(
1228 stride_o_dim,
1229 stride_o_token,
1230 # others
1228 - pad_slot_id,
1231 + null_block_id,
1232 # META
1233 HAS_BIAS=bias is not None,
1234 KERNEL_WIDTH=width,
@@ -1234,7 +1237,7 @@ def causal_conv1d_update(
1237 IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
1238 IS_SPEC_DECODING=num_accepted_tokens is not None,
1239 NP2_STATELEN=np2_statelen,
1237 - USE_PAD_SLOT=pad_slot_id is not None,
1240 + HAS_NULL_BLOCK=null_block_id is not None,
1241 BLOCK_N=256,
1242 )
1243 if unsqueeze:
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+16 -16
@@ -10,7 +10,7 @@ from packaging import version
10 from vllm import _custom_ops as ops
11 from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
12 from vllm.triton_utils import HAS_TRITON, tl, triton
13 -from vllm.v1.attention.backends.utils import PAD_SLOT_ID
13 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
14
15 TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
16
@@ -75,7 +75,7 @@ def _selective_scan_update_kernel(
75 out_ptr,
76 state_batch_indices_ptr,
77 dst_state_batch_indices_ptr,
78 - pad_slot_id,
78 + null_block_id,
79 num_accepted_tokens_ptr,
80 cu_seqlens_ptr,
81 # Matrix dimensions
@@ -203,7 +203,7 @@ def _selective_scan_update_kernel(
203
204 mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)
205 if HAS_STATE_BATCH_INDICES:
206 - mask &= state_batch_idx != pad_slot_id
206 + mask &= state_batch_idx != null_block_id
207 state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32)
208
209 if HAS_DT_BIAS:
@@ -257,7 +257,7 @@ def _selective_scan_update_kernel(
257 if IS_SPEC_DECODING:
258 dst_idx_ptr = dst_state_batch_indices_ptr + i_t * stride_dst_state_indices_T
259 token_dst_idx = tl.load(dst_idx_ptr).to(tl.int64)
260 - if token_dst_idx != pad_slot_id:
260 + if token_dst_idx != null_block_id:
261 token_dst_ptrs = (
262 state_ptr_base
263 + token_dst_idx * stride_state_batch
@@ -329,7 +329,7 @@ def selective_state_update(
329 dt_softplus=False,
330 state_batch_indices=None,
331 dst_state_batch_indices=None,
332 - pad_slot_id=PAD_SLOT_ID,
332 + null_block_id=NULL_BLOCK_ID,
333 out=None,
334 num_accepted_tokens=None,
335 cu_seqlens=None,
@@ -348,12 +348,12 @@ def selective_state_update(
348 D: (dim,) or (nheads, dim)
349 z: (batch, dim) or (batch, nheads, dim)
350 dt_bias: (dim,) or (nheads, dim)
351 - pad_slot_id: int
352 - if cache_indices is passed, lets the kernel identify padded
353 - entries that will not be processed,
354 - for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
355 - in this case, the kernel will not process entries at
356 - indices 0 and 3
351 + null_block_id: int
352 + if state_batch_indices is passed, lets the kernel identify
353 + padded entries that will not be processed,
354 + for example: state_batch_indices = [null_block_id, 1, 20,
355 + null_block_id] in this case, the kernel will not process
356 + entries at indices 0 and 3
357 out: Preallocated ssm output tensor. Assume same shape as x.
358 In-place updated.
359 num_accepted_tokens: (batch,)
@@ -488,7 +488,7 @@ def selective_state_update(
488 out,
489 state_batch_indices,
490 dst_state_batch_indices,
491 - pad_slot_id,
491 + null_block_id,
492 num_accepted_tokens,
493 cu_seqlens,
494 N,
@@ -550,7 +550,7 @@ def selective_scan_fn(
550 query_start_loc=None,
551 cache_indices=None,
552 has_initial_state=None,
553 - pad_slot_id=PAD_SLOT_ID,
553 + null_block_id=NULL_BLOCK_ID,
554 block_size=1024,
555 block_idx_first_scheduled_token=None,
556 block_idx_last_scheduled_token=None,
@@ -588,10 +588,10 @@ def selective_scan_fn(
588 indicate if the ssm_state at the corresponding index should be
589 used as initial state. Not providing argument assumes
590 there's no initial state
591 - pad_slot_id: int
591 + null_block_id: int
592 if cache_indices is passed, lets the kernel identify padding entries
593 that will not be processed,
594 - for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
594 + for example: cache_indices = [null_block_id, 1 ,20 ,null_block_id]
595 in this case, the kernel will not process entries at indices 0 and 3
596 block_size: int
597 The block size to align the cached states to
@@ -643,7 +643,7 @@ def selective_scan_fn(
643 cache_indices,
644 has_initial_state,
645 ssm_states,
646 - pad_slot_id,
646 + null_block_id,
647 block_size,
648 block_idx_first_scheduled_token,
649 block_idx_last_scheduled_token,
vllm/v1/attention/backends/gdn_attn.py
+3 -3
@@ -14,7 +14,7 @@ from vllm.v1.attention.backend import (
14 CommonAttentionMetadata,
15 )
16 from vllm.v1.attention.backends.utils import (
17 - PAD_SLOT_ID,
17 + NULL_BLOCK_ID,
18 compute_causal_conv1d_metadata,
19 mamba_get_block_table_tensor,
20 split_decodes_and_prefills,
@@ -341,7 +341,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
341 spec_state_indices_tensor, non_blocking=True
342 )
343 spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size]
344 - spec_state_indices_tensor[num_spec_decodes:].fill_(PAD_SLOT_ID)
344 + spec_state_indices_tensor[num_spec_decodes:].fill_(NULL_BLOCK_ID)
345
346 self.spec_sequence_masks[:num_spec_decodes].copy_(
347 spec_sequence_masks[:num_spec_decodes], non_blocking=True
@@ -387,7 +387,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
387 non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[
388 :batch_size
389 ]
390 - non_spec_state_indices_tensor[num_decodes:].fill_(PAD_SLOT_ID)
390 + non_spec_state_indices_tensor[num_decodes:].fill_(NULL_BLOCK_ID)
391
392 self.non_spec_query_start_loc[: num_decodes + 1].copy_(
393 non_spec_query_start_loc, non_blocking=True
vllm/v1/attention/backends/mamba_attn.py
+2 -2
@@ -15,7 +15,7 @@ from vllm.v1.attention.backend import (
15 CommonAttentionMetadata,
16 )
17 from vllm.v1.attention.backends.utils import (
18 - PAD_SLOT_ID,
18 + NULL_BLOCK_ID,
19 compute_causal_conv1d_metadata,
20 mamba_get_block_table_tensor,
21 split_decodes_and_prefills,
@@ -504,7 +504,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
504 state_indices_tensor_d, non_blocking=True
505 )
506 state_indices_tensor_d = self.state_indices_tensor_d[:padded_bs]
507 - state_indices_tensor_d[metadata.num_decodes :] = PAD_SLOT_ID
507 + state_indices_tensor_d[metadata.num_decodes :] = NULL_BLOCK_ID
508
509 if self.use_spec_decode and num_accepted_tokens is not None:
510 assert query_start_loc_d is not None
vllm/v1/attention/backends/mla/indexer.py
-6
@@ -366,12 +366,6 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
366 seq_lens = common_attn_metadata.seq_lens[:num_decodes]
367 block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
368
369 - # Padded CUDA graph requests have block_table entries of -1.
370 - # Clamp to 0 to prevent OOB access in the DeepGEMM kernel.
371 - # This is safe because padded requests have seq_lens=0, so the
372 - # kernel produces no meaningful output for those rows.
373 - block_table.clamp_(min=0)
374 -
369 max_decode_len = int(decode_lens_cpu.max().item())
370 next_n = 1 + self.num_speculative_tokens
371 use_native = not self.use_flattening and max_decode_len == next_n
vllm/v1/attention/backends/utils.py
+1
@@ -42,6 +42,7 @@ KVCacheLayoutType = Literal["NHD", "HND"]
42 _KV_CACHE_LAYOUT_OVERRIDE: KVCacheLayoutType | None = None
43
44 PAD_SLOT_ID = -1
45 +NULL_BLOCK_ID = 0
46
47
48 def is_valid_kv_cache_layout(value: str) -> bool:
vllm/v1/worker/gpu_model_runner.py
+4 -3
@@ -122,6 +122,7 @@ from vllm.v1.attention.backend import (
122 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
123 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
124 from vllm.v1.attention.backends.utils import (
125 + NULL_BLOCK_ID,
126 create_fast_prefill_custom_backend,
127 get_dcp_local_seq_lens,
128 reorder_batch_to_split_decodes_and_prefills,
@@ -2135,9 +2136,9 @@ class GPUModelRunner(
2136 blk_table = self.input_batch.block_table[kv_cache_gid]
2137 blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded)
2138
2138 - # Fill unused with -1. Needed for reshape_and_cache in full cuda
2139 - # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID
2140 - blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)
2139 + # Fill unused block table entries with NULL_BLOCK_ID (null block)
2140 + # for CUDAGraph padding. Block 0 is reserved for padding.
2141 + blk_table_tensor[num_reqs:num_reqs_padded].fill_(NULL_BLOCK_ID)
2142 return blk_table_tensor
2143
2144 assert slot_mappings is not None