diff --git a/deepinfer/neuron_ops/nki_attention.py b/deepinfer/neuron_ops/nki_attention.py
index 7e526d4..66f1be2 100644
--- a/deepinfer/neuron_ops/nki_attention.py
+++ b/deepinfer/neuron_ops/nki_attention.py
@@ -255,7 +255,9 @@ def flash_mla_decode(
         if mask is not None:
             invalid = invalid | mask
         scopes.append((cache, torch.where(invalid, torch.zeros_like(idx), idx), invalid))
-    if rows == PREFILL_BUCKET:
+    if rows == PREFILL_BUCKET and any(idx.shape[1] >= _SHARED_SCOPE_MIN for _, idx, _ in scopes):
+        # Only a dense compressed-KV scope needs the shared-gather path. Top-k and SWA scopes are per row either way,
+        # so the kernel takes them as in decode (the torch path was 41.6 ms a layer for a 128-token chunk, bench4).
         out = _chunk_attention(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)
     else:
         out = sparse_decode(q.reshape(rows, h_q, d), scopes, softmax_scale, attn_sink)