Skip to content

Make auto-CP planning stable across head shards - #32

Draft
Micuks wants to merge 2 commits into
QwenLM:mainfrom
Micuks:fix/global-head-autocp-planning
Draft

Make auto-CP planning stable across head shards#32
Micuks wants to merge 2 commits into
QwenLM:mainfrom
Micuks:fix/global-head-autocp-planning

Conversation

@Micuks

@Micuks Micuks commented Jul 29, 2026

Copy link
Copy Markdown

What changed

Auto-CP planning can now use a logical head count from
FLASH_QLA_AUTO_CP_PLANNING_HEADS. Kernel calls still use their local head
count. The override only affects the occupancy estimate and segment plan.

The default behavior is unchanged when the variable is unset. Invalid values
raise a clear error, and the logical head count must be a positive multiple of
the local head count.

This PR also adds a standalone regression and records the H800 reproduction.
The test calls FlashQLA directly and does not depend on a model repository.

Why

A logical H16 operation may be issued as four H4 calls after head sharding.
The current scheduler plans each call from its local head count, so the H16 and
H4 calls can choose different segment boundaries.

For the reproduced mixed-length input, H16 used 18 auto-CP segments while H4
used 34. Their BF16 outputs differed. Disabling auto-CP made the same inputs
bit-exact, which isolates the difference to the auto-CP path.

Setting FLASH_QLA_AUTO_CP_PLANNING_HEADS=16 gives both launch shapes the same
18-segment plan and restores bit-exact output while keeping auto-CP enabled.

Validation

The standalone test ran on an NVIDIA H800 with PyTorch 2.8.0+cu128:

Case H16/H4 segments Max abs Exact fraction
Without override, auto-CP on 18 / 34 0.00048828125 0.9808532
Without override, auto-CP off n/a 0 1.0
Fixed, auto-CP on 18 / 18 0 1.0

I also ran:

python3 -m py_compile tests/validate_global_head_autocp_planning.py
git diff --check

The regression covers one mixed-length BF16 forward case on SM90. It does not
claim parity for every shape, GPU architecture, backward path, or full
training run.

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