| """Per-wave bulk-training context for varlen ring and sequence-parallel canaries.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
|
|
| _varlen_cu_seqlens: torch.Tensor | None = None |
| _varlen_max_seqlen: torch.Tensor | None = None |
|
|
|
|
| def set_bulk_varlen_context_boundary( |
| *, |
| cu_seqlens: torch.Tensor | None, |
| max_seqlen: torch.Tensor | None, |
| ) -> None: |
| """Install varlen metadata for the active bulk CUDA wave.""" |
|
|
| global _varlen_cu_seqlens, _varlen_max_seqlen |
| _varlen_cu_seqlens = cu_seqlens |
| _varlen_max_seqlen = max_seqlen |
|
|
|
|
| def clear_bulk_varlen_context_boundary() -> None: |
| """Drop varlen metadata after one wave completes.""" |
|
|
| global _varlen_cu_seqlens, _varlen_max_seqlen |
| _varlen_cu_seqlens = None |
| _varlen_max_seqlen = None |
|
|
|
|
| def bulk_varlen_context_boundary() -> tuple[torch.Tensor | None, torch.Tensor | None]: |
| """Return active varlen ``(cu_seqlens, max_seqlen)`` when installed.""" |
|
|
| return _varlen_cu_seqlens, _varlen_max_seqlen |
|
|
