Skip to content

fix: force TileLang-contiguous inputs for Mamba3 MIMO - #994

Open
Chessing234 wants to merge 4 commits into
state-spaces:mainfrom
Chessing234:fix/mamba3-mimo-seqlen1-contiguous
Open

fix: force TileLang-contiguous inputs for Mamba3 MIMO#994
Chessing234 wants to merge 4 commits into
state-spaces:mainfrom
Chessing234:fix/mamba3-mimo-seqlen1-contiguous

Conversation

@Chessing234

@Chessing234 Chessing234 commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Force physical last-stride=1 after .contiguous() before TileLang MIMO kernels.
  • Fixes forward crash at seqlen=1 where PyTorch still reports contiguous despite stride(-1) != 1.

Fixes #985

Test plan

  • Mamba3(..., is_mimo=True)(x) with x.shape == (B, 1, D) on CUDA

@Chessing234
Chessing234 force-pushed the fix/mamba3-mimo-seqlen1-contiguous branch 2 times, most recently from 6ee2e0c to 8976fb0 Compare July 18, 2026 15:01
PyTorch treats size-1 trailing dims as contiguous even when stride(-1) != 1,
which breaks TileLang kernels at seqlen=1.
@Chessing234
Chessing234 force-pushed the fix/mamba3-mimo-seqlen1-contiguous branch from 8976fb0 to bb34615 Compare July 18, 2026 15:01
Keeps the state-spaces#985 stride fix importable without pulling Triton kernels,
so CPU tests can cover size-1 trailing dims.
Asserts physical last-stride=1 for size-1 trailing dims without CUDA.
Avoids requiring Triton at collection time for the CPU stride helper
tests.
@Chessing234

Copy link
Copy Markdown
Contributor Author

Follow-up commits on this branch:

  1. Moved ensure_tilelang_contiguous into tilelang_layout.py so it can be tested without importing Triton.
  2. CPU unit tests for the size-1 trailing-stride case.
  3. CUDA smoke: Mamba3(is_mimo=True) with seqlen=1 in tests/modules/test_mamba3_varlen.py.

Still targeting #985. Ready for another look when you have a minute.

@Chessing234
Chessing234 force-pushed the fix/mamba3-mimo-seqlen1-contiguous branch from 545d063 to 208be9f Compare August 5, 2026 10:24
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.

Mamba3 MIMO forward doesn't work with seq_length=1

1 participant