File size: 987 Bytes
919fd68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """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
|