From 2c44356cc5b114d512c3a0ea5f9e59cea928dc41 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Tue, 15 Sep 2026 06:12:15 +0000 Subject: [PATCH] fix(iluvatar): correct HD128 paged prefill dispatch Use the existing InfiniOps global warp kernel for Iluvatar HD128 prefill while preserving other backends dispatch. Fold query-token counts above 65535 into grid.z to support large batches. Add ChatGLM3 and MiniCPM4 regression coverage, HD128 smoke coverage, and grid-boundary tests. Both model issues share the same prefill fix. --- .../paged_attention_prefill_infinilm/kernel.h | 14 +++++- tests/conftest.py | 17 +++++-- tests/test_flash_attn_varlen_func.py | 46 +++++++++++++++++++ 3 files changed, 71 insertions(+), 6 deletions(-) diff --git a/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h b/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h index a1b2a10e2..655b6158e 100644 --- a/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h +++ b/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(num_heads_), @@ -117,6 +120,15 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm { } else { dim3 legacy_grid(static_cast(num_heads_), static_cast(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((total_q_tokens_ + 65534) / 65535); + } + } op::paged_attention_prefill::cuda:: PagedAttentionPrefillWarpGlobalKernel diff --git a/tests/conftest.py b/tests/conftest.py index 407851500..4abaf3d09 100644 --- a/tests/conftest.py +++ b/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 ) diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index c8a8badf4..ec9d76a40 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/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