Skip to content

fix(attention): Scope FMHA cubin loading by mask type - #148

Open
suharvest wants to merge 1 commit into
NVIDIA:mainfrom
suharvest:pr/fix-fmha-mask-scoped-loading
Open

fix(attention): Scope FMHA cubin loading by mask type#148
suharvest wants to merge 1 commit into
NVIDIA:mainfrom
suharvest:pr/fix-fmha-mask-scoped-loading

Conversation

@suharvest

Copy link
Copy Markdown

What does this PR do?

Scopes context-FMHA cubin loading and caching by the requested attention mask
type.

Previously the load cache key contained only data type and SM. Constructing a
list then attempted to load every matching data-type/SM cubin before the mask
type was considered during function lookup. An unrelated custom-mask cubin
could therefore fail with CUDA_ERROR_INVALID_IMAGE while initializing a
causal model.

The requested mask type now participates in:

  • the loader cache key and equality/hash;
  • metadata filtering before cubin loading;
  • availability and dispatch lookup.

Every caller passes its mask explicitly; the public load function has no
default mask argument. A real custom-mask request still loads the custom-mask
cubin and propagates load failures—there is no fallback to a different mask.

Scope

  • Context-FMHA loader/cache implementation and declaration.
  • Necessary LLM, ViT, and existing unit-test call sites.
  • No cubin, kernel, model, or fallback changes.

Validation

  • Source-contract checks verify mask-aware equality, hashing, metadata
    filtering, cache lookup, availability, and dispatch.
  • All nine declarations/calls use the explicit three-argument API.
  • The loader does not catch or suppress custom-mask load exceptions.
  • git diff --check and merge-tree validation pass.

The validation machine has no CUDA/TensorRT toolchain, so native compilation,
the non-custom-mask positive engine case, and a real invalid custom-mask cubin
negative test are still requested from CI/device testing.

Closes #143.

Signed-off-by: suharvest <suharvest@gmail.com>
@suharvest
suharvest requested a review from a team July 25, 2026 12:08
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.

[Bug][Runtime]: context FMHA loads unrelated mask-type cubins and can fail with CUDA_ERROR_INVALID_IMAGE

1 participant