Skip to content

[Fix] Honor strides in vector activations - #1369

Merged
yzhangcs merged 2 commits into
fla-org:mainfrom
adenzhou1350:fix/strided-vector-activations
Oct 11, 2026
Merged

yzhangcs merged 2 commits into
fla-org:mainfrom
adenzhou1350:fix/strided-vector-activations

Conversation

@adenzhou1350

Copy link
Copy Markdown
Contributor

Summary

Fix incorrect activation values and gradients for one-dimensional strided views, such as x[::3]. The layout guard previously accepted every rank-one tensor even though these kernels use unit column offsets; the existing contiguous fallback now also handles strided vector inputs and upstream gradients for sigmoid, logsigmoid, swish, SwiGLU, SigmoidGLU and PowGLU.

Contiguous vectors and the higher-dimensional row-stride path are unchanged. This follows the non-contiguous support in #765; it does not extend scalar support or change the separate Ascend implementation.

Test plan

  • RTX 5090, Python 3.12, PyTorch 2.11.0+cu130, Triton 3.6.0, default Triton backend, FLA_CI_ENV=0.
  • New FP32 cases cross contiguous/stride-3 inputs with contiguous/stride-3 incoming gradients for all six functions, checking outputs and every supported gradient against the existing differentiable references at the existing 1e-3 tolerance.
  • On baseline f3bec720f7847d4f176872b34f002cc580e2cc6e, the identical regression tests give 6 passes and 19 failures (18 numerical failures plus the layout-helper assertion). On this head all 25 pass. Independent public-call probes also reproduce the numerical failures before the fix and pass after it.
  • FLA_CI_ENV=0 pytest tests/modules/test_activation.py -q --tb=short: 103 cases passed across two bounded runs. The first completed 72 cases before a 900-second cold-compilation limit; the next deselected exactly those 72 and passed the remaining 31. This is not a single uninterrupted full-suite pass. No cases failed or skipped; the second run reported three PyTorch/pytest setup advisories, not tolerance warnings.
  • pre-commit run --files fla/modules/activations/ops.py tests/modules/test_activation.py: passed.
  • Ran python scripts/find_dependent_tests.py fla/modules/activations/ops.py and audited shared callers. Its broad transitive layer/model test list was not run; this change is limited to rank-one layout admission, and the activation suite covers the existing dense/chunked and compiled paths. No full-model or non-NVIDIA validation claim.

Benchmark / NCU

N/A for kernel benchmarks: no kernel code or higher-rank dispatch changed. This restores correctness through the existing copy fallback for previously misread vector views, so the incorrect no-copy baseline is not a valid performance comparison for those views. No performance-neutrality, throughput or speedup claim.

Breaking changes

None. Public signatures, dtypes and tolerances are unchanged.

Checklist

  • I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
  • I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
  • Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
  • Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
  • This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.
  • I understand this PR exceeds 500 changed lines and have explained its single purpose and why it should stay together below.

@yzhangcs
yzhangcs merged commit 363df60 into fla-org:main Oct 11, 2026
14 of 15 checks passed
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.

2 participants