Skip to content

Handle all-empty varlen batches - #32

Open
0z5a wants to merge 2 commits into
MoonshotAI:masterfrom
0z5a:codex/zero-length-varlen
Open

Handle all-empty varlen batches#32
0z5a wants to merge 2 commits into
MoonshotAI:masterfrom
0z5a:codex/zero-length-varlen

Conversation

@0z5a

@0z5a 0z5a commented Aug 17, 2026

Copy link
Copy Markdown

Summary

  • Define zero-length sequence semantics for cu_seqlens with repeated offsets.
  • Avoid constructing invalid zero-extent TMA descriptors when the entire varlen batch is empty.
  • Preserve initial_state in final_state, or write zeros when no initial state is supplied.

Dependency

This PR is stacked on #29. Until #29 merges, the diff includes its tile-prefix commit; after #29 lands, this PR can be rebased to contain only the empty-batch fix.

Reproduction

On the #29 baseline, mixed-empty batches (leading, middle, or consecutive empty sequences) already pass, while T_total == 0 aborts in the C++ extension when zero-extent TMA descriptors are constructed.

Validation

RTX 5080 (SM120): python -m pytest tests/test_empty_varlen.py -vv — 32 passed, covering mixed/all-empty batches, H=1/4, BF16/FP32 state, and calls with or without initial_state.

@0z5a 0z5a changed the title Support-empty-varlen-batches Handle all-empty varlen batches Aug 17, 2026
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