[Bugfix][Qwen4Exp] fix state index strides in fused PLE conv (#55375)
Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
Harry Huang committed
Sep 5, 2026 at 22:02 UTC
28e605fb330093fee5a2b98aebc19d9c268c1822
2 files changed
+31
-3
tests/models/qwen4_exp/test_ple.py
+23
-1
@@ -1077,6 +1077,8 @@ class _ConvBatchCase:
1077
dilation: int = 3
1078
spec_query_len: int = 1
1079
graph_padding: int = 0
1080
+ state_index_stride: int = 1
1081
+ include_null_state: bool = True
1082
1083
1084
def _make_conv_metadata(
@@ -1123,9 +1125,18 @@ def _make_conv_metadata(
1125
dtype=torch.int32,
1126
device=device,
1127
)
1128
+ if case.state_index_stride > 1:
1129
+ strided_indices = torch.full(
1130
+ (non_spec_state_indices.numel(), case.state_index_stride),
1131
+ NULL_BLOCK_ID,
1132
+ dtype=torch.int32,
1133
+ device=device,
1134
+ )
1135
+ strided_indices[:, 0] = non_spec_state_indices
1136
+ non_spec_state_indices = strided_indices[:, 0]
1137
if case.num_decodes > 1:
1138
non_spec_state_indices[case.num_decodes - 1] = NULL_BLOCK_ID
1128
- if num_prefills > 1:
1139
+ if num_prefills > 1 and case.include_null_state:
1140
non_spec_state_indices[case.num_decodes + 1] = NULL_BLOCK_ID
1141
1142
spec_query_start_loc = torch.tensor(
@@ -1203,6 +1214,15 @@ def _make_conv_metadata(
1214
_ConvBatchCase(prefill_query_lens=(37, 0, 5, 128)),
1215
id="prefill",
1216
),
1217
+ pytest.param(
1218
+ _ConvBatchCase(
1219
+ prefill_query_lens=(16, 676),
1220
+ channels=10240,
1221
+ state_index_stride=4,
1222
+ include_null_state=False,
1223
+ ),
1224
+ id="prefill-strided-state-indices",
1225
+ ),
1226
pytest.param(
1227
_ConvBatchCase(
1228
spec_query_lens=(1, 4, 4, 0),
@@ -1243,6 +1263,8 @@ def test_fused_conv_correctness(
1263
) -> None:
1264
device = torch.device("cuda")
1265
metadata, num_real_tokens = _make_conv_metadata(case, device)
1266
+ if case.state_index_stride > 1:
1267
+ assert metadata.state_indices_tensor.stride(0) == case.state_index_stride
1268
module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
1269
nn.Module.__init__(module)
1270
module.conv_state_len = (case.kernel_size - 1) * case.dilation
vllm/models/qwen4_exp/nvidia/ops/ple.py
+8
-2
@@ -330,6 +330,7 @@ def _ple_conv_kernel(
330
has_token_map,
331
num_reqs,
332
bs_iters,
333
+ state_idx_stride,
334
state_bs,
335
state_ws,
336
state_cs,
@@ -386,7 +387,7 @@ def _ple_conv_kernel(
387
else:
388
slot_off = tl.full([], 0, tl.int32)
389
389
- sid = tl.load(state_idx_ptr + r).to(tl.int64)
390
+ sid = tl.load(state_idx_ptr + r * state_idx_stride).to(tl.int64)
391
state_ok = sid != NULL_STATE_ID
392
sid_safe = tl.where(state_ok, sid, 0)
393
if HAS_INIT:
@@ -484,6 +485,7 @@ def _ple_conv_writeback_kernel(
485
has_init_ptr,
486
token_idx_ptr,
487
has_token_map,
488
+ state_idx_stride,
489
state_bs,
490
state_ws,
491
state_cs,
@@ -501,7 +503,7 @@ def _ple_conv_writeback_kernel(
503
c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
504
c_mask = c_offs < C
505
504
- sid = tl.load(state_idx_ptr + r).to(tl.int64)
506
+ sid = tl.load(state_idx_ptr + r * state_idx_stride).to(tl.int64)
507
state_ok = sid != NULL_STATE_ID
508
if not state_ok:
509
return
@@ -612,6 +614,8 @@ def ple_conv(
614
615
num_warps = 4 if mode == "prefill" else 8
616
launch_pdl = current_platform.is_arch_support_pdl()
617
+ # Pure-prefill indices can be a strided block-table column view.
618
+ state_idx_stride = state_indices.stride(0)
619
620
# Constexpr flags eliminate accesses to optional None arguments. Without a
621
# token map, state_indices is an unused but device-resident placeholder.
@@ -628,6 +632,7 @@ def ple_conv(
632
token_indices is not None,
633
num_reqs,
634
binary_search_iters,
635
+ state_idx_stride,
636
state_bs,
637
state_ws,
638
state_cs,
@@ -654,6 +659,7 @@ def ple_conv(
659
has_initial_states,
660
token_indices if token_indices is not None else state_indices,
661
token_indices is not None,
662
+ state_idx_stride,
663
state_bs,
664
state_ws,
665
state_cs,