Skip to content

fix(zamba): add use_associative_scan config flag to avoid torch.compile slowdown - #48331

Open
msnliu wants to merge 1 commit into
huggingface:mainfrom
msnliu:fix/zamba-use-associative-scan
Open

fix(zamba): add use_associative_scan config flag to avoid torch.compile slowdown#48331
msnliu wants to merge 1 commit into
huggingface:mainfrom
msnliu:fix/zamba-use-associative-scan

Conversation

@msnliu

@msnliu msnliu commented Aug 26, 2026

Copy link
Copy Markdown

CPU CI GPU run-slow

ZambaMambaMixer.forward() passes use_associative_scan=False directly to mamba_selective_scan, which falls back to a sequential Python loop that makes torch.compile slow.

The other Mamba-family models (mamba, jamba, falcon_mamba) instead read this setting from their config and forward it to the scan function. This PR adds the same plumbing to ZambaConfig and ZambaMambaMixer.

Changes:

  1. configuration_zamba.py: add use_associative_scan: bool = True field
  2. modeling_zamba.py: read config.use_associative_scan in init and pass it in forward() instead of the hardcoded False

@vasqu

@github-actions

Copy link
Copy Markdown
Contributor

[For maintainers] Suggested jobs to run (before merge)

run-slow: zamba

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 32953204500
Result: failure | Grafana metrics are not available yet.

Code quality check failed: test jobs were skipped. Fix the code quality issues and push again to run tests.

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.

1 participant