fix(iluvatar): correct HD128 paged prefill dispatch

src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h

@@ -63,7 +63,10 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
             }
           }();
 
-          if constexpr (kHeadSize == 128) {
+          // The HD128 pipeline produces incorrect results on Iluvatar.
+          // Keep its native prefill on the global warp kernel.
+          if constexpr (kHeadSize == 128 &&
+                        Backend::kDeviceType != Device::Type::kIluvatar) {
             if (block_size_ == 256) {
               constexpr int kWarps = 8;
               dim3 pipe_grid(static_cast<unsigned>(num_heads_),
@@ -117,6 +120,15 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
           } else {
             dim3 legacy_grid(static_cast<unsigned>(num_heads_),
                              static_cast<unsigned>(total_q_tokens_));
+            if constexpr (Backend::kDeviceType == Device::Type::kIluvatar) {
+              // The kernel folds grid.z into its token index when grid.y
+              // would exceed CUDA's 65535 limit (e.g. B16 x I4096).
+              if (total_q_tokens_ > 65535) {
+                legacy_grid.y = 65535;
+                legacy_grid.z =
+                    static_cast<unsigned>((total_q_tokens_ + 65534) / 65535);
+              }
+            }
             op::paged_attention_prefill::cuda::
                 PagedAttentionPrefillWarpGlobalKernel<Backend::kDeviceType,
                                                       TIndex, TData, kHeadSize>

tests/conftest.py

@@ -499,16 +499,23 @@ def _is_smoke_cutlass_scaled_mm_case(params):
 def _is_smoke_flash_attn_varlen_func_case(params):
     if params.get("device") == "cuda" and params.get("implementation_index") == 0:
         return (
-            params.get("q_lens") == (2, 3)
-            and params.get("k_lens") == (130, 300)
-            and params.get("num_heads") == 4
+            (
+                params.get("q_lens"),
+                params.get("k_lens"),
+                params.get("num_heads"),
+                params.get("head_dim"),
+                params.get("use_alibi"),
+            )
+            in (
+                ((2, 3), (130, 300), 4, 64, True),
+                ((2, 3), (130, 300), 4, 128, True),
+                ((13,), (13,), 32, 128, False),
+            )
             and params.get("num_kv_heads") == 2
             and params.get("causal") is True
             and params.get("window_size") == (-1, -1)
             and params.get("scale") is None
             and params.get("paged") is True
-            and params.get("use_alibi") is True
-            and params.get("head_dim") == 64
             and params.get("dtype") == torch.float16
         )
 

tests/test_flash_attn_varlen_func.py

@@ -24,6 +24,9 @@
         ((5, 2), (3, 6), 4, 2, True, (-1, -1), 0.125, False, False),
         ((4, 3), (6, 2), 4, 2, False, (2, 1), None, False, False),
         ((4, 3), (6, 2), 4, 2, True, (2, 1), None, False, False),
+        ((13,), (13,), 32, 2, True, (-1, -1), None, True, False),
+        # MiniCPM4 batched prefill; catches the Iluvatar HD128 pipeline regression.
+        ((1024,) * 16, (1024,) * 16, 32, 2, True, (-1, -1), None, True, False),
         ((2, 3), (130, 300), 4, 2, True, (-1, -1), None, True, True),
     ),
 )
@@ -288,6 +291,49 @@ def test_ascend_paged_prefill_follows_nontrivial_block_table(
     torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
 
 
+@pytest.mark.parametrize("total_tokens", (65535, 65536, 131071))
+def test_iluvatar_prefill_grid_limit(device, implementation_index, total_tokens):
+    if device != "cuda" or implementation_index != 0:
+        pytest.skip("this launch geometry regression targets Iluvatar native")
+    if 0 not in infini.ops.FlashAttnVarlenFunc.active_implementation_indices(
+        "iluvatar"
+    ):
+        pytest.skip("requires the Iluvatar native provider")
+
+    # Each one-token sequence attends to the same physical cache slot. This
+    # exercises grid.y/z boundaries without quadratic reference allocations.
+    q = torch.zeros((total_tokens, 1, 128), dtype=torch.float16, device=device)
+    k = torch.zeros((1, 256, 1, 128), dtype=q.dtype, device=device)
+    v = torch.randn_like(k)
+    cumulative = torch.arange(total_tokens + 1, dtype=torch.int32, device=device)
+    block_table = torch.zeros((total_tokens, 1), dtype=torch.int32, device=device)
+    out = torch.full_like(q, math.nan)
+    infini.ops.flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cumulative,
+        cumulative,
+        None,
+        block_table,
+        1,
+        1,
+        0.0,
+        None,
+        True,
+        (-1, -1),
+        0.0,
+        False,
+        False,
+        out,
+        None,
+        None,
+        stream=get_stream(q.device),
+        implementation_index=implementation_index,
+    )
+    torch.testing.assert_close(out, v[0, 0].expand_as(out), rtol=2e-3, atol=2e-3)
+
+
 def test_flash_attn_varlen_func_non_default_stream(device, implementation_index):
     if device == "cuda":
         accelerator = torch.cuda