[Longxmas]

Outcome

Variable-length DSpark/DFlash BSHD training with FA4 repeatedly compiled custom attention kernels when the context length or anchor count changed. The same geometry already reused its kernel within one process.

There were two independent cache-key problems:

1. The mask factory captured concrete context/query/key lengths and anchor counts, so each new geometry changed the modifier hash used by FA4's forward and backward caches.
2. Nested CuTe helper functions contributed process-specific function representations to the closure hash. Consequently, enabling persistent caching alone still missed custom attention kernels in a new process.

This change moves row geometry to runtime auxiliary metadata and bounds, and gives the modifier a stable hash covering the complete factory source and every remaining semantic specialization. It also forwards FA4's opt-in persistent-cache settings to Ray actors while retaining node-local cache paths for static multi-node workers. FA4's dtype/head/tile/broadcast and singleton-stride ABI distinctions remain intact.

Reproduction and before/after comparison

Environment: an SM90 machine, Python 3.12.13, PyTorch 2.13.0+cu130, CUDA 13.0, FA4 4.0.0b34, CUTLASS DSL 4.7.1; compiler and ptxas optimization level 3. Both versions used the same GPU and synthetic BF16 inputs, with persistent caching enabled. Each version started with separate empty caches, then reused its cache in a second process.

Geometry: batch 1, Q/KV heads 4/1, head dimension 128, block size 7, sparse FA4. Each process runs (context length, anchor blocks) = (512,32), (512,32), (640,32), (640,32), (512,32), (640,40), (640,40).

Forward + backward call	Baseline	Fixed
First process, initial (512,32)	12.764 s; 5 compile calls	13.384 s; 5 compile calls
Repeat (512,32) in the same process	5.109 ms; 0 compile calls	5.198 ms; 0 compile calls
New context (640,32)	11.216 s; 2 compile calls	3.622 ms; 0 compile calls
New block count (640,40)	11.378 s; 2 compile calls	4.239 ms; 0 compile calls
Second process, first (512,32)	11.921 s; 2 compile calls	201.428 ms; 0 compile calls

The seven-call sequences invoked cute.compile 9 / 6 times in baseline processes one/two, versus 5 / 0 times after the fix. Modifier hashes now match across lengths and processes.

These are operator-level compilation/cache diagnostics. Wall time covers synchronized forward and backward; compilation counts and CPU time are measured by wrapping cute.compile. All observed outputs and gradients were finite. End-to-end training speedup is not measured here.

Reproduce the regression

From the PR branch, run the same tests against the unchanged baseline production code and the fix:

git fetch https://github.com/inclusionAI/ASpec.git main
git worktree add --detach ../aspec-fa4-before "$(git merge-base HEAD FETCH_HEAD)"
cp tests/test_direct_cp.py ../aspec-fa4-before/tests/test_direct_cp.py

# Baseline: expected cache-reuse assertion failure.
(cd ../aspec-fa4-before && CUDA_VISIBLE_DEVICES=0 ASPEC_RUN_GPU_INTEGRATION=1 \
  python aspec/_environment.py python -m pytest -q tests/test_direct_cp.py -k variable_lengths)

# PR branch: all four parameterized cases must pass.
CUDA_VISIBLE_DEVICES=0 ASPEC_RUN_GPU_INTEGRATION=1 \
  python aspec/_environment.py python -m pytest -q tests/test_direct_cp.py -k variable_lengths

Verification

• Baseline regression reproduced: the same new GPU test, copied into the baseline checkout without changing production files, fails with row geometry recompiled an existing FA4 dispatch family.
• Fixed GPU regression: 4 passed. Compares output, LSE, and gradients with mathematical attention across variable lengths, changing anchors, masked blocks, truncated tails, context/query windows, and causal drafts. Shared test caches also guard against aliasing distinct mask semantics. Legitimate singleton-stride dispatch families remain separate.
• Two-GPU Direct CP: 1 passed on each rank. Forward and gradients match the unsharded mathematical reference.
• Consolidated CPU/source checks: 32 passed, 5 skipped. Covers Direct CP, environment forwarding, actor defaults, and source quality after consolidating the tests. The broader related suite also passed (162 passed, 11 skipped, 6 subtests). Skips are not counted as GPU evidence.
• Full repository Ruff check, Python compileall, maintained Markdown link check, and git diff --check: passed.

Runtime validation is on SM90 with the versions above. Other GPU architectures and end-to-end training/serving performance are not measured by this PR.

Compatibility and provenance

Model parameters, attention masks, losses, data selection, checkpoint formats, and default persistent-cache enablement are unchanged. The per-forward context-boundary tensor uses the existing metadata-cache lifetime. Persistent reuse remains explicitly enabled with FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 and FLASH_ATTENTION_CUTE_DSL_CACHE_DIR.

This modifies ASpec's retained NVIDIA/Megatron-derived adapter; existing source headers and license attribution are preserved. No dependency changes or additional upstream implementation copies are introduced.

[zhouaihui]

LGTM