[Bugfix][Qwen4Exp] fix state index strides in fused PLE conv (#55375)

Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>

tests/models/qwen4_exp/test_ple.py

@@ -1077,6 +1077,8 @@ class _ConvBatchCase:
     dilation: int = 3
     spec_query_len: int = 1
     graph_padding: int = 0
+    state_index_stride: int = 1
+    include_null_state: bool = True
 
 
 def _make_conv_metadata(
@@ -1123,9 +1125,18 @@ def _make_conv_metadata(
         dtype=torch.int32,
         device=device,
     )
+    if case.state_index_stride > 1:
+        strided_indices = torch.full(
+            (non_spec_state_indices.numel(), case.state_index_stride),
+            NULL_BLOCK_ID,
+            dtype=torch.int32,
+            device=device,
+        )
+        strided_indices[:, 0] = non_spec_state_indices
+        non_spec_state_indices = strided_indices[:, 0]
     if case.num_decodes > 1:
         non_spec_state_indices[case.num_decodes - 1] = NULL_BLOCK_ID
-    if num_prefills > 1:
+    if num_prefills > 1 and case.include_null_state:
         non_spec_state_indices[case.num_decodes + 1] = NULL_BLOCK_ID
 
     spec_query_start_loc = torch.tensor(
@@ -1203,6 +1214,15 @@ def _make_conv_metadata(
             _ConvBatchCase(prefill_query_lens=(37, 0, 5, 128)),
             id="prefill",
         ),
+        pytest.param(
+            _ConvBatchCase(
+                prefill_query_lens=(16, 676),
+                channels=10240,
+                state_index_stride=4,
+                include_null_state=False,
+            ),
+            id="prefill-strided-state-indices",
+        ),
         pytest.param(
             _ConvBatchCase(
                 spec_query_lens=(1, 4, 4, 0),
@@ -1243,6 +1263,8 @@ def test_fused_conv_correctness(
 ) -> None:
     device = torch.device("cuda")
     metadata, num_real_tokens = _make_conv_metadata(case, device)
+    if case.state_index_stride > 1:
+        assert metadata.state_indices_tensor.stride(0) == case.state_index_stride
     module = Qwen4ExpPLELayer.__new__(Qwen4ExpPLELayer)
     nn.Module.__init__(module)
     module.conv_state_len = (case.kernel_size - 1) * case.dilation

vllm/models/qwen4_exp/nvidia/ops/ple.py

@@ -330,6 +330,7 @@ def _ple_conv_kernel(
     has_token_map,
     num_reqs,
     bs_iters,
+    state_idx_stride,
     state_bs,
     state_ws,
     state_cs,
@@ -386,7 +387,7 @@ def _ple_conv_kernel(
         else:
             slot_off = tl.full([], 0, tl.int32)
 
-    sid = tl.load(state_idx_ptr + r).to(tl.int64)
+    sid = tl.load(state_idx_ptr + r * state_idx_stride).to(tl.int64)
     state_ok = sid != NULL_STATE_ID
     sid_safe = tl.where(state_ok, sid, 0)
     if HAS_INIT:
@@ -484,6 +485,7 @@ def _ple_conv_writeback_kernel(
     has_init_ptr,
     token_idx_ptr,
     has_token_map,
+    state_idx_stride,
     state_bs,
     state_ws,
     state_cs,
@@ -501,7 +503,7 @@ def _ple_conv_writeback_kernel(
     c_offs = pid_c * BLOCK_C + tl.arange(0, BLOCK_C)
     c_mask = c_offs < C
 
-    sid = tl.load(state_idx_ptr + r).to(tl.int64)
+    sid = tl.load(state_idx_ptr + r * state_idx_stride).to(tl.int64)
     state_ok = sid != NULL_STATE_ID
     if not state_ok:
         return
@@ -612,6 +614,8 @@ def ple_conv(
 
     num_warps = 4 if mode == "prefill" else 8
     launch_pdl = current_platform.is_arch_support_pdl()
+    # Pure-prefill indices can be a strided block-table column view.
+    state_idx_stride = state_indices.stride(0)
 
     # Constexpr flags eliminate accesses to optional None arguments. Without a
     # token map, state_indices is an unused but device-resident placeholder.
@@ -628,6 +632,7 @@ def ple_conv(
         token_indices is not None,
         num_reqs,
         binary_search_iters,
+        state_idx_stride,
         state_bs,
         state_ws,
         state_cs,
@@ -654,6 +659,7 @@ def ple_conv(
             has_initial_states,
             token_indices if token_indices is not None else state_indices,
             token_indices is not None,
+            state_idx_stride,
             state_bs,
             state_ws,
             state_cs,