Repository navigation
[WIP]feat(turbo): route attention through Primus-Turbo's flex_attention when enabled - #1033
Open
botaohu001 wants to merge 16 commits into
Open
botaohu001 wants to merge 16 commits into
botaohu001 wants to merge 16 commits into
Conversation
Route Megatron core attention through the Primus-Turbo flex_attention compat layer instead of calling flash_attn_func directly, behind a new config key. The compat layer takes a torch-FlexAttention block_mask/score_mod, classifies it, and dispatches the recognized patterns onto the same Turbo kernels, so with the default causal configuration this is numerically identical to the direct call (same kernel, same arguments). What it buys is the ability to express masks and score modifications that flash_attn_func has no argument for, without leaving Turbo's fast path. Config schema (primus/configs/modules/megatron/primus_turbo.yaml): use_turbo_flex_attention: false # new, off by default turbo_flex_attention_mask_mod: null # optional "module:attr" hook turbo_flex_attention_score_mod: null # optional "module:attr" hook Implementation: TurboFlexAttention mirrors flash_attn_func's call signature and is substituted for self.attn / self.attn_func in PrimusTurboAttention and PrimusTurboLocalAttention, so neither forward() changes. It maps causal / window_size onto BlockMasks cached per (pattern, shape, device) -- the compat layer memoizes its classification by object identity, so reuse keeps the per-call overhead at zero. Dispatch prefers flex_attention_bshd (layout-native, no transposes); older Turbo builds fall back to the bhsd entry with an explicit round-trip. Unsupported combinations raise instead of degrading to the direct call, so an unsupported run fails at model-build time rather than training on silently by the old path: context parallelism, enable_turbo_attention_float8, deterministic_mode, packed (THD) sequences, return_attn_probs, a sliding window without a causal mask, and use_turbo_flex_attention without use_turbo_attention. Adds 33 CPU unit tests (no GPU, no real Turbo build required) and a performance-tuning doc section.
…cache Three gaps found while auditing the switch for real training use. 1. Sink dtype. PrimusTurboAttention allocates self.sinks as a bfloat16 Parameter (mirroring gpt-oss; flash_attn_func takes it as-is), but the compat layer validates sink.dtype == torch.float32 and raises otherwise. use_sink_attention + use_turbo_flex_attention therefore died with a ValueError on the first forward -- i.e. the switch was not a drop-in for the gpt-oss configs. Cast in the adapter with a differentiable .float(), so the bf16 Parameter still receives a gradient. 2. Build-time rejections. deterministic_mode was only caught on the first forward, after model init and data loading; reset_attention_mask was not caught at all, and would have silently dropped per-document boundaries inside a packed sample (the Turbo path receives only the causal flag). Both now raise from build_turbo_flex_attention, next to the existing context-parallel and fp8 rejections. 3. BlockMask cache growth. The cache was an unbounded dict keyed by (pattern, q_len, kv_len, device). Fixed-length pretraining holds one or two entries, but a variable-length workload adds one BlockMask per distinct length for the lifetime of the process. Bounded LRU (64 entries) with a one-shot RuntimeWarning on first eviction, so a pathological shape distribution is reported rather than leaked. Tests 33 -> 43 (CPU-only, no GPU and no real Turbo build required). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…_varlen qkv_format="thd" previously had no path at all on the Turbo attention side. PrimusTurboAttention.forward read qkv_format out of packed_seq_params but only permuted for sbhd/bhsd, and never forwarded cu_seqlens to self.attn -- so a packed batch reached flash_attn_func with the document boundaries dropped and tokens attending across documents, with nothing raised anywhere. TurboFlexAttention now accepts cu_seqlens_q/cu_seqlens_kv/max_seqlen_q/max_seqlen_kv and dispatches 3D input to flex_attention_varlen, which applies document-internal causal masking. The packed path builds no BlockMask (the boundaries are given, not inferred), so it neither probes nor touches the LRU cache. mask_mod, score_mod and bias are refused there rather than silently dropped. The extra kwargs are gated on _attn_accepts_packed_seq, set only when the flex switch is on, so the direct flash_attn_func binding is byte-for-byte unchanged; on that path qkv_format="thd" now raises instead of silently losing the boundaries. THD output is reshaped to [total_tokens, H*D] (the generic 3-arg view was a no-op on 3D). reset_attention_mask stays rejected: it encodes boundaries in a dense sample's mask tensor, which is a different mechanism with no argument to carry it. Tests 43 -> 56 (new TestPackedTHD, 12 cases); pre-commit green.
botaohu001
requested review from
Xiaoming-AMD,
limou102 and
wenxie-amd
as code owners
August 25, 2026 10:34
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| if len(_BLOCK_MASK_CACHE) > _BLOCK_MASK_CACHE_MAX: | ||
| _BLOCK_MASK_CACHE.popitem(last=False) | ||
| if not _EVICTION_WARNED: | ||
| _EVICTION_WARNED = True |
Co-authored-by: Cursor <cursoragent@cursor.com> # Conflicts: # docs/04-technical-guides/performance-tuning.md
…section) The file was deleted while resolving a merge conflict, which broke the Read the Docs (Sphinx) build: docs/sphinx/_toc.yml.in and 8+ in-repo links still reference it. Restore it as a clean 3-way merge of main's #976 path updates and this PR's added Primus-Turbo flex_attention tuning section. Co-authored-by: Cursor <cursoragent@cursor.com>
wenxie-amd
previously approved these changes
Aug 25, 2026
…ig SWA on flex End-to-end coverage work (EXP17) turned up three parameters that PrimusTurboAttention accepted and then ignored. Each one runs clean, reports no error, and trains a different model than the config asked for: * attention_bias is a forward parameter that nothing reads. Megatron delivers ALiBi as an attention bias, and the kernel call passes alibi_slopes=None unconditionally, so an ALiBi config would train without ALiBi. Now raises. * window_size[1] was dropped by the dense path while window_size[0] was applied, so (4, 4) was silently treated as (4, 0). Now raises; the packed path forwards both bounds to the varlen backend verbatim and is unaffected. * A user mask_mod replaced any requested window instead of composing with it. The combination is reachable from a plain config (turbo_flex_attention_mask_mod plus sink_sliding_window), so it is rejected at model-build time as well as in the call. Also narrows the config.window_size assert. It read "primus_turbo does not support sliding window attention", but the sink path was already driving the same aiter kernel with a window. The flex path expresses a window as a mask_mod, so it can honour config.window_size; the direct flash_attn_func binding has no window argument and still rejects it. config.window_size together with sink_sliding_window is refused -- two sources for one quantity, with different per-layer rules. Verified end to end on one MI300 (20 iters, seq 2048, 4 layers, seed 1234): config window_size=(512,0) matches a hand-written sliding_window(512) mask_mod bit-for-bit across all 20 iterations, and the same config on the non-flex path exits at build time with the new message. Unit suite 55 -> 61 tests, all passing.
Primus-Turbo now hands a bhsd view straight to the backend instead of materialising it, so the fallback through the torch-layout entry no longer costs four copies per forward on current builds. The call contract is unchanged; only the comments were wrong. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…mask PrimusTurboAttention.forward and PrimusTurboLocalAttention.forward both accept an attention_mask tensor and never read it. That is fine for the default mask -- create_attention_mask_in_dataloader defaults to true, and the tensor it produces is exactly torch.tril(...), i.e. the causal flag already forwarded to the kernel. It is not fine under reset_attention_mask=true, where the dataset additionally zeroes the cross-document blocks (gpt_dataset.py::_get_ltor_masks_and_position_ids). No Turbo entry point takes a mask tensor, so those boundaries were dropped and tokens attended across documents with nothing raised anywhere -- a silent correctness bug, not a crash. Reject the configuration at build time on both direct paths, via one shared helper so the reasoning lives in a single place. The flex path already raised for itself with a message naming the THD alternative, so the helper defers to it. Packed sequences (qkv_format='thd') remain the supported way to get per-document masking: there the boundaries arrive as cu_seqlens and are forwarded explicitly. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ntry, and un-gate window_size Three silent-drop / over-restriction fixes on the Megatron extension. deterministic (flex path): TurboFlexAttention hard-coded deterministic=False, so the flag could never reach the kernel. It is now threaded through both the dense and the packed call, guarded by an inspect.signature capability probe (_entry_accepts) so an older Primus-Turbo whose flex entries have no such parameter still raises instead of dropping it. build_turbo_flex_attention's deterministic_mode rejection is likewise conditional on the probe. THD packing (direct path): PrimusTurboAttention now binds flash_attn_varlen_func alongside the dense entry and forwards cu_seqlens / max_seqlen positionally, so qkv_format="thd" works without the flex layer. Bound only at context_parallel_size == 1: flash_attn_varlen_usp_func wants global cu_seqlens while each rank holds a local shard, and guessing would reintroduce cross-document attention. reset_attention_mask is deliberately still rejected -- it yields a mask tensor, not cu_seqlens. window_size: the blanket config.window_size rejection is now fp8-only. forward() already threaded window_size into every attention call, and flash_attn_func / flash_attn_varlen_func / flash_attn_usp_func all honour it (usp in the backward too); only the fp8 kernel asserts window == (-1,-1). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
self.attn_kwargs carries the context-parallel process groups (ulysses_group / ring_group) and is only populated when cp_size > 1; self.attn_varlen is only bound when cp_size == 1. So today the two are mutually exclusive and the direct varlen branch never drops anything -- but that is an invariant held by two separate conditions in __init__, invisible from the call site. Assert it here so that relaxing either gate later fails loudly instead of silently running attention over the local shard only, which would quietly train a different model. No behaviour change on any path reachable today.
PrimusTurboAttention.forward() passed `sink=sink_tensor` unconditionally,
but only flash_attn_func has that parameter. flash_attn_usp_func,
flash_attn_fp8_usp_func, flash_attn_fp8_func and flash_attn_varlen_func
do not, so every context_parallel_size > 1 run died on a bare
TypeError: flash_attn_usp_func() got an unexpected keyword argument 'sink'
regardless of whether a sink was configured -- sink=None is still passing
the kwarg -- and with nothing in the message to say why. fp8 hit the same
wall at CP == 1. Reproduced on 2x MI355 and confirmed present at the
merge-base, so this is not something the flex branch introduced.
Add turbo_flex_attention.sink_kwargs_for(entry, sink, *, where), which
answers the question three ways:
* entry takes one -> forward it (None included; that is its default)
* entry does not, no sink -> omit the kwarg. Nothing is lost.
* entry does not, sink set -> raise. Dropping it would change the softmax
denominator of every query and silently
train a different model.
The middle case is the fix; the third is the reason it is a helper and not
a try/except. The dense and packed paths probe their own entry separately,
since the two answers genuinely differ.
It lives in turbo_flex_attention.py rather than at the call site because
that module has no module-level megatron import and is therefore unit
testable; primus_turbo.py is not.
Also sharpen the CP == 1-only comment on the attn_varlen binding. The
kernel half of that hedge is no longer in doubt: measured on 2x MI355,
flash_attn_varlen_usp_func is bit-identical to the single-card varlen
kernel given a local token shard plus global cu_seqlens, including with a
document straddling the shard boundary. What remains unverifiable is
whether Megatron's PackedSeqParams carries global or already-sharded
boundaries under CP, so the rejection stays -- the comment now says that
instead of implying the kernel is the problem.
Tested: 7 new unit tests for the three branches; 79 passed in
test_turbo_flex_attention.py (1 pre-existing failure, unrelated: it
assumes a Primus-Turbo build whose flex entries lack `deterministic`).
The two call sites inside PrimusTurboAttention cannot be exercised here --
no megatron on this machine -- so they are changed against measured entry
signatures rather than end-to-end execution.
…eachable
PrimusTurboAttention extends te.pytorch.DotProductAttention, not Megatron's
TEDotProductAttention -- but forward() read self.kept_packed_seq_params, which
only the latter defines, in its own __init__ (Megatron-LM
megatron/core/extensions/transformer_engine.py:1309; checked in-container that
TE 2.15's DotProductAttention does not define it either).
So every call carrying packed_seq_params died on
AttributeError: 'PrimusTurboAttention' object has no attribute
'kept_packed_seq_params'
before reaching any of the packing logic, and the error said nothing about
packing. That is upstream of the attn_varlen binding (1c296eb), the THD x CP
rejection (aeb0fcc) and the varlen sink gate (c162da9): the entire THD path
has never run. It only surfaced now because the Megatron-LM submodule was
uninitialised, so the adapter could not be exercised end to end at all.
The dict comprehension was only ever read for one key, so there is nothing to
reconstruct: read qkv_format off the dataclass. A qkv_format of None now falls
back to the module default instead of propagating -- None matches neither "sbhd"
nor "thd", so it used to skip the layout permute entirely and hand the kernel a
transposed tensor.
Second fix, same theme: cu_seqlens_q_padded / cu_seqlens_kv_padded were read off
PackedSeqParams by nobody. They describe a different packing from cu_seqlens_*
-- the latter are real token counts, the former include per-sequence padding.
TE takes both and honours the distinction; the varlen entry bound here takes one
pair, so the unpadded boundaries would be applied to a padded tensor, attending
over padding and misplacing every document after the first. Refuse loudly.
Verified on MI355 (gfx950) with the Megatron-LM submodule initialised at its
pinned SHA, driving the real module (gpu_runs_mi355/exp23_adapter_e2e.py):
dense no-sink, dense with-sink (output moves by 2.95, so the sink reaches the
kernel rather than being dropped), packed THD, packed+sink refused, and at
world=2 cp=2 the CP dense path plus a counterfactual showing the pre-c162da9f
behaviour still raises TypeError.
…ection The build stopped rejecting deterministic_mode unconditionally in 1c296eb: the compat layer threads 'deterministic' through now, so the build only rejects it on a Turbo whose flex entries predate the parameter. That commit added both halves of the new contract to TestBuild -- accepted on a current turbo, rejected on an older one -- but left the old unconditional test in TestBuildTimeRejections behind. The stub backend's _Recorder.__call__ takes **kwargs, which _entry_accepts (correctly) reads as accepting 'deterministic', so the build does not raise and the stale pytest.raises reports DID NOT RAISE. The failure was evidence that the production code is right, not that it is wrong. Remove it and say in the class docstring where the real contract is tested. No production code touched. 79 passed, 0 failed.
…g them
PrimusTurboLocalAttention.forward accepted both and read neither. The kernel
call is hardcoded to bias=None / alibi_slopes=None and there is no thd branch
at all, so:
- attention_bias is how Megatron delivers ALiBi. A user setting
position_embedding_type=alibi never types the word "bias" and gets a model
trained without ALiBi, silently.
- packed_seq_params carries the document boundaries. Dropped, a packed batch
is attended to as one contiguous sequence and every document sees every
earlier one. Wrong numbers, wrong gradients, clean log.
Raise NotImplementedError for both, matching the wording and reasoning of the
attention_bias refusal already in PrimusTurboAttention.forward. The packed
refusal points at that spec, which does implement qkv_format='thd'.
Adds 7 CPU tests. They call forward unbound against a stub, so no GPU, no
process groups, no get_args. Besides the refusals they pin that the message
names ALiBi, that an empty packed_seq_params is still refused (emptiness is not
permission), and that the ordinary dense call is byte-for-byte unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Xiaoming-AMD
approved these changes
Sep 1, 2026
botaohu001
changed the title
This branch has not been deployed
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
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.
Adds an opt-in path that routes Megatron attention through Primus-Turbo's
flex_attentioncompat layer, so a torch-FlexAttention-style call dispatches toTurbo's tuned aiter kernels instead of generated Triton.Off by default (
use_turbo_flex_attention: false); nothing changes unless it is