Port Kunlun adaptations to kunlun_v2#7
Merged
Conversation
- Remove the varlen ban in flash2 checks so the sequence-packed two_way attention path can select the flash2 backend (I4_ATTN_BACKENDS=flash2) instead of falling back to torch_sdpa. - Normalize the varlen LSE layout: Kunlun's flash_attn_varlen_func returns [batch, H, max_seqlen] (padded per sample) whereas NVIDIA returns [H, total_tokens]. Un-pad using cu_seqlens_q so both builds produce [H, total_tokens]. - Add explicit `import torch` for the un-padding torch.cat.
* Replace NVML device affinity with XPUML * Migrate remaining pynvml usages to XPUML Replace NVML calls in conftest, device utils, inference args, and the device monitor callback with torch_xmlir XPUML bindings.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Port the three Kunlun-specific commits from
kunlunonto the cleankunlun_v2branch based on currentmain.This PR does not merge
kunlunor include its previous main-sync history.Ported commits
1997b87feda56e78c6958179c05494fb534ac772->dad930b(changes for cosmos3 super inference)c555de4f371954937993c576a5464fee45940d84->9f04cd3(Enable flashAttn2 varlen backend on Kunlun P800)4b037c5562bc743c0f3aba08a70986aad46d42fc->a26345b(Replace NVML device affinity with XPUML)The commits were cherry-picked in the requested order. The only cleanup was removal of trailing whitespace already present in the first source commit; it was folded into the corresponding ported commit.
Scope
kunlun_v2692c410583a2e1271cdc355ffe2ffe574c2ae6254923586and allkunlunmerge/revert historyValidation
git diff --check origin/kunlun_v2...HEAD: passedpython -m compileall -q conftest.py cosmos_framework: passedpre-commitworkflow: passedpytest -q: stopped during collection because the local environment is missingomegaconfruff/uvis unavailable in the local environment