Skip to content

fix(iluvatar): correct HD128 paged prefill dispatch - #990

Open
gongchensu wants to merge 1 commit into
InfiniTensor:masterfrom
gongchensu:fix/iluvatar-prefill-head128
Open

gongchensu wants to merge 1 commit into
InfiniTensor:masterfrom
gongchensu:fix/iluvatar-prefill-head128

Conversation

@gongchensu

@gongchensu gongchensu commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

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.

Summary

  • Keep Iluvatar HD128 paged prefill on the existing global warp kernel in src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h, avoiding the incorrect HD128 pipeline path on Iluvatar.
  • Preserve the existing HD128 pipeline dispatch for other backends.
  • Fold total query-token counts above 65535 into grid.z for Iluvatar so large batch prefill launches stay within CUDA grid limits.
  • Add ChatGLM3 and MiniCPM4 HD128 regression shapes to tests/test_flash_attn_varlen_func.py.
  • Add HD128 native prefill coverage to the smoke selection in tests/conftest.py.
  • Add launch-boundary tests for total query-token counts 65535, 65536, and 131071.

Motivation

The Iluvatar HD128 pipeline path produced incorrect paged prefill results for ChatGLM3 and MiniCPM4. For the ChatGLM3-shaped HD128 case, the original path had a maximum error around 4.611; using the existing global warp kernel reduces it to about 0.000967.

Large batch prefill also placed the total query-token count directly in grid.y, so workloads such as batch 16 with input length 4096 exceeded CUDA's 65535 grid.y limit. The global warp kernel already supports folding the extra tokens into grid.z; this PR enables that launch geometry only for Iluvatar.

No linked issue is available yet; reproducers and automated regression coverage are included in this PR.

Type of Change

  • feat — new feature / new operator / new platform
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

paste smoke test output here

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA
Iluvatar
MetaX
Cambricon
Moore
Ascend
Full `pytest` output (optional)
paste here

Benchmark / Performance Impact

Notes for Reviewers

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.
@gongchensu
gongchensu requested a review from a team September 18, 2026 02:31
@gongchensu gongchensu self-assigned this Sep 18, 2026
@gongchensu
gongchensu requested a review from voltjia September 18, 2026 02:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant