Skip to content

refactor(runtime)!: adopt the modern Infini stack - #506

Draft
voltjia wants to merge 33 commits into
mainfrom
refactor/adopt-modern-infini-stack
Draft

refactor(runtime)!: adopt the modern Infini stack#506
voltjia wants to merge 33 commits into
mainfrom
refactor/adopt-modern-infini-stack

Conversation

@voltjia

@voltjia voltjia commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Rebase the full InfiniCore-to-InfiniLM migration onto InfiniLM main at 80bb09ecebc9aabf198b9b866a89456bca1df946.
  • Preserve the original migration while synchronizing the effective InfiniCore changes that landed after the migration source diverged.
  • Keep runtime context, tensor/graph integration, operator adapters, distributed wrappers, and Python bindings owned by InfiniLM.
  • Re-audit the migration against current InfiniRT and InfiniOps master, then unlock the NVIDIA capabilities whose complete provider and InfiniLM execution paths are now validated.
  • Keep incomplete surfaces gated with errors that describe the current integration gap instead of the stale generic "kernels are unavailable" reason.

The refreshed branch contains 32 commits on top of current main; the final head is 8c2d3f5acc845481d4aff5767881ec0d5d1bc1c9.

Related to InfiniTensor/InfiniCore#1373.

Rebase and InfiniCore delta audit

The original migration source ended at InfiniCore d616d4f032e65f901287cbf4e1e6ec4b551703f9. That source and the audited InfiniCore line diverge at 1590de7b0b469af45af5e073996d5a3fab277041; the migration synchronization endpoint was 9855c0c6227093dc35b5a49c3b24792e9d99171e.

The post-migration audit produced four effective synchronization groups:

  1. Segmented graph replay and recordable point-to-point communication, while retaining InfiniLM's stronger owner-scoped cancellation and runtime/allocation lifetime handling.
  2. RoPE cache accessors, degenerate transposed Parameter loading, and improved tensor-view diagnostics.
  3. Graph-aware, device-native Tensor::ones and Tensor::zeros implemented through canonical Fill/InfiniRT paths.
  4. Paged-cache 32768-token chunking, SwiGLU row-boundary chunking, and static-attention value-head-width fixes.

No fifth effective synchronization group was found after excluding changes already migrated, superseded by stronger InfiniLM behavior, or unrelated to this migration boundary.

Current upstream capability audit

The component contract still follows InfiniCore #1406:

  • InfiniCore #1406 head: 55cfe5e6761c4ebb5e8eb77b65301479a7b6032c
  • InfiniRT: 0cdbb16967e15f2e055dea1ec9641617bf3b6cf6
  • InfiniOps: f890afb4b2327f13ccdd3c1b6b0d49567c5fe00d
  • InfiniCCL: e8ccc0cb23ca5b1b1d63be29ba61ea031de72807

The current upstream heads were also audited:

  • InfiniRT master: 6b256e0ab0853d9e2a7ef059bac22c847bd02251. Its only delta from the pin is a Mars header fix.
  • InfiniOps master: 2d05aca88dc12a79a6aec26caef5bf2251752978. Its deltas from the pin are Ascend fixes, CMake package export, and native example repair; none changes the NVIDIA operator/API surface used here.

No dependency pin bump is needed for this refresh. The selected InfiniOps closure now contains 24 operators, including Add plus the linked FlashAttention prefill/decode providers and paged-cache write provider.

Newly unlocked

Explicit FlashAttention

  • --attn=flash-attn now selects the existing FLASH_ATTN backend instead of throwing.
  • The provider chain uses linked flash_attn_varlen_func and flash_attn_with_kvcache, plus reshape_and_cache_flash.
  • Construction rejects unsupported combinations before rank workers start: non-NVIDIA devices, non-paged caches, zero blocks, block sizes not divisible by 256, non-FP16/BF16 dtypes, and head dimensions outside the provider contract.

Linear bias

  • Dense and packed linear paths now execute Gemm with beta=0 and compose broadcast Add for bias.
  • Row-parallel bias is added once after all-reduce; column-parallel bias remains sharded.
  • Replicated parameter metadata no longer uses an invalid negative TP dimension.
  • Fused pre-transposed weights split along the physical transposed dimension.

Dense model factory

The modern backend now compiles and enables:

baichuan, chatglm, fm9g, fm9g7b, glm4, internlm3, llama, minicpm, minicpm4, qwen2, and qwen3.

The public MiniCPM4 checkpoint keeps its raw model_type=minicpm; the Python reader intentionally normalizes it to minicpm4, which is now selected, linked, and covered by a cross-layer contract.

Pre-transposed weights

--pre-transpose is not Ascend-only. Its CLI description now reflects the validated one-time pre-transposition and packed-Gemm behavior on NVIDIA.

Gates intentionally retained

  • FlashInfer attention: the selected linked provider implements sampling only, not an attention backend.
  • Compressed-tensors W8A8: InfiniLM lacks the end-to-end activation quantization and scaled-Gemm integration.
  • AWQ/GPTQ: current InfiniOps exposes repack helpers, but InfiniLM has no dense AWQ or GPTQ/Marlin Gemm execution path.
  • MXFP4/Quark: no InfiniOps-backed InfiniLM execution path is selected.
  • INT8 KV cache: selected attention providers consume FP16/BF16 KV tensors and InfiniLM has no INT8 conversion/consumption chain.
  • GPT-2: NVIDIA LayerNorm has no registered InfiniOps provider.
  • Mistral: the current implementation does not consume sliding_window; a short prompt completing is not sufficient semantic coverage.
  • Mamba, Qwen3-Next, and Qwen3.5: current InfiniOps master has no NVIDIA causal_conv1d, selective-scan, or gated-delta-rule provider chain required by these models.
  • MoE and multimodal families: complete model-specific operator/adapter paths are not present in the selected closure.
  • Other devices: this refresh validates NVIDIA A100 only.

Validation

Repository and native build

  • Full static contract suite: 74/74 passed.
  • Build-script unit suite: 19/19 passed.
  • Ruff check/format, clang-format dry-run/Werror, and git diff --check: passed.
  • Focused InfiniOps FlashAttention provider tests on A100: 35 passed, 1 skipped (two-GPU guard on a single-GPU run).
  • Latest InfiniOps 2d05aca8 configured, built, and installed with the complete 24-operator manifest.
  • Final 17-file InfiniLM snapshot: forced full NVIDIA build and install passed.
  • In-place broadcast Add [3, 7] + [7]: output aliases the destination and max_abs_error=0.

Real-weight dense models

Two-token NVIDIA eager smokes passed for:

  • Llama-3.2-3B-Instruct
  • FM9G 9g-8B
  • Baichuan2-7B-Chat
  • ChatGLM3-6B
  • InternLM3-8B-Instruct
  • GLM-4-9B
  • MiniCPM4-8B through the normalized minicpm4 implementation
  • Qwen2-compatible FM9G-70B under TP8
  • Qwen3-0.6B

The fm9g7b and native minicpm keys share the already compiled FM9G registrar; they did not receive separate checkpoint smokes.

FlashAttention

  • Qwen3-0.6B explicit FlashAttention: eager and segmented graph passed.
  • Qwen3-32B BF16 TP4 explicit FlashAttention: passed.
  • Llama-3.2-3B explicit FlashAttention: eager and segmented graph passed.

Bias, TP, and pre-transposition

A Qwen3-0.6B overlay enabled nonzero attention Q/K/V bias, attention-output bias, and MLP gate/up/down bias:

  • Static raw forward: zero/nonzero bias under TP1 and TP2 all passed.
  • Zero-to-nonzero logits effects were consistent between TP1 and TP2; the remaining TP1/TP2 BF16 differences matched the zero-bias accumulation-order baseline.
  • Column-parallel bias shards and row-parallel replicated bias shapes matched the intended TP contract.
  • Pre-transposed TP1/TP2 raw forward passed with the expected packed weight shapes.
  • Paged explicit FlashAttention: zero/nonzero bias x TP1/TP2 all passed with identical generated tokens.
  • Nonzero bias + pre-transposition + paged FlashAttention passed under TP1/TP2 eager and TP1 segmented graph capture/replay.

Preserved migration baseline

  • InfiniRT Release: 12/12 CTest tests passed.
  • InfiniCCL two-GPU all-reduce produced the correct result.
  • Wheel install/import, Torch-preloaded linkage, GPU copy, graph cancel/re-record/replay, Fill, and RoPE smokes passed.
  • Qwen3-0.6B static/paged, eager/graph, batch-2, greedy/non-greedy, TP2, PP2, and TP2+PP2 coverage remains valid.

Type of change

  • refactor
  • fix
  • test
  • docs
  • build / CI
  • breaking change

Landing order

  1. feat(linked): add FlashInfer sampling provider InfiniOps#930, #931, #935, and #945: merged.
  2. fix(build): regenerate communication bridge for header changes InfiniCCL#57, [BUG] test_ppl.py run fail #58, and Feature/use logsoft max in ppl #59: open.
  3. refactor!: reduce InfiniCore to a component manifest InfiniCore#1406: open.
  4. This PR.

This PR remains draft until the remaining upstream dependencies land.

@voltjia
voltjia force-pushed the refactor/adopt-modern-infini-stack branch 3 times, most recently from 7c36e2f to 077867b Compare August 13, 2026 15:26
voltjia added 25 commits August 25, 2026 12:53
BREAKING CHANGE: legacy InfiniCore runtime and inference APIs are removed in favor of InfiniRT, InfiniOps, and the migrated InfiniLM runtime.
Carry segmented P2P graph replay, graph-aware tensor initializers, runtime compatibility helpers, and long-context chunking into the InfiniLM-owned runtime.

Preserve the stronger migrated graph ownership, cancellation, and allocation lease semantics while reconciling current InfiniLM attention behavior.
@voltjia
voltjia force-pushed the refactor/adopt-modern-infini-stack branch from 9439ea6 to 9d64e14 Compare August 25, 2026 05:12
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