[Bugfix][Model] Fix Inkling NVIDIA sconv cache block-size mismatch

tests/models/inkling/test_sconv_cache_layout.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""NVIDIA runtime conv-block-size tracking (see the ROCm twin test).
+
+The KV-cache planner enlarges the conv group's block size (W=4 -> e.g. 8)
+when the attention page is a multiple of the conv page
+(``unify_kv_cache_spec_page_size``). Slot mappings are then built with the
+enlarged size, so NVIDIA call sites must index with the *runtime* block size,
+not the kernel window size.
+"""
+
+import torch
+from torch import nn
+
+from vllm.models.inkling.nvidia.sconv_swa_attn import InklingConvState
+
+
+def test_runtime_sconv_block_size_tracks_unified_cache_page():
+    owner = InklingConvState.__new__(InklingConvState)
+    nn.Module.__init__(owner)
+    owner.block_size = 4
+
+    owner.kv_cache = torch.tensor([])
+    assert owner.cache_block_size == 4
+
+    owner.kv_cache = torch.empty(2, 1, 32, 1024)
+    assert owner.cache_block_size == 32

vllm/models/inkling/nvidia/attention.py

@@ -257,7 +257,7 @@ def forward(
                 fa_md.slot_mapping,
                 off_k,
                 off_v,
-                self.conv_owner.block_size,
+                self.conv_owner.cache_block_size,
                 log_scaling if not self.is_local else None,
             )
             q = q.view(num_tokens, self.num_heads, self.head_dim)

vllm/models/inkling/nvidia/model.py

@@ -108,7 +108,7 @@ def _sconv_add_norm(
             m.slot_mapping,
             off_s,
             ws,
-            sconv.owner.block_size,
+            sconv.owner.cache_block_size,
             shared_tensor=shared_delta,
         )
 

vllm/models/inkling/nvidia/sconv_swa_attn.py

@@ -176,9 +176,10 @@ def __init__(
         self.num_kv_heads = num_kv_heads // tp_size
         hidden_per_head = hidden_size // num_kv_heads
         # Packed per-head width: K + V + attn-output chunk + mlp-output chunk,
-        # padded to a power of two so every layer's conv page is the same size
-        # and an exact multiple of the attention page (the page unifier then
-        # scales attention block sizes instead of padding).
+        # padded to a power of two so every layer's conv page is the same size.
+        # The page unifier may still scale the logical conv block size up to
+        # match a larger attention page; index the bound cache with its runtime
+        # token dim (``cache_block_size``), not this kernel window size.
         raw_head_size = 2 * head_dim + 2 * hidden_per_head
         self.head_size = 1 << (raw_head_size - 1).bit_length()
         self.sliding_window = kernel_size
@@ -205,6 +206,13 @@ def __init__(
 
     def forward(self): ...
 
+    @property
+    def cache_block_size(self) -> int:
+        """Return the block size used by cache metadata and physical indexing."""
+        if self.kv_cache.numel() > 0:
+            return self.kv_cache.shape[2]
+        return self.block_size
+
     def get_attn_backend(self) -> type[AttentionBackend]:
         return InklingSconvBackend
 

vllm/models/inkling/nvidia/short_conv.py

@@ -72,7 +72,13 @@ def forward(self, x: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
             return x
 
         off_s, ws = self.owner.stream_ranges[self.stream_idx]
-        block_size = self.owner.block_size
+        # The hybrid KV-cache planner can enlarge the logical conv block so
+        # that its physical page size matches the attention caches (for
+        # example, W=4 becomes 8 when the attention page doubles the conv
+        # page). Metadata slot mappings are built with that enlarged size, so
+        # index the bound cache with its runtime token dimension rather than
+        # the kernel window size.
+        block_size = self.owner.cache_block_size
         x = x.contiguous()
         weight = self.weight.squeeze(1)  # (dim, W)