Purge usage of pytorch named tensors (#25911)

python/sglang/srt/debug_utils/comparator/aligner/axis_aligner.py

@@ -12,6 +12,7 @@ from sglang.srt.debug_utils.comparator.dims_spec import (
     DimSpec,
     _SingletonDimUtil,
     parse_dims,
+    without_dim_names,
 )
 from sglang.srt.debug_utils.comparator.log_sink import log_sink
 from sglang.srt.debug_utils.comparator.utils import Pair, _FrozenBase
@@ -213,6 +214,6 @@ def execute_axis_aligner_plan(
     pattern: Optional[str] = plan.pattern.x if side == "x" else plan.pattern.y
 
     if pattern is not None:
-        tensor = rearrange(tensor.rename(None), pattern)
+        tensor = rearrange(without_dim_names(tensor), pattern)
 
     return tensor

python/sglang/srt/debug_utils/comparator/aligner/reorderer/executor.py

@@ -8,8 +8,10 @@ from sglang.srt.debug_utils.comparator.aligner.reorderer.types import (
     ZigzagToNaturalThdParams,
 )
 from sglang.srt.debug_utils.comparator.dims_spec import (
+    apply_dim_names,
+    get_dim_names,
     resolve_dim_by_name,
-    strip_dim_names,
+    without_dim_names,
 )
 
 
@@ -47,8 +49,8 @@ def _reorder_zigzag_to_natural_thd(
     Each seq in seq_lens is independently reordered from zigzag to natural order
     along the given dim.
     """
-    stripped: torch.Tensor = strip_dim_names(tensor)
-    names: tuple[Optional[str], ...] = tensor.names
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    stripped: torch.Tensor = without_dim_names(tensor)
 
     split_sizes: list[int] = list(seq_lens)
     remainder: int = stripped.shape[dim] - sum(split_sizes)
@@ -74,7 +76,7 @@ def _reorder_zigzag_to_natural_thd(
     result: torch.Tensor = torch.cat(reordered_segments, dim=dim)
 
     if names[0] is not None:
-        result = result.refine_names(*names)
+        result = apply_dim_names(result, list(names))
     return result
 
 
@@ -86,8 +88,8 @@ def _reorder_zigzag_to_natural(
     Generalized from Megatron-LM _undo_attention_load_balancing
     (megatron/core/ssm/mamba_context_parallel.py:360-373).
     """
-    stripped: torch.Tensor = strip_dim_names(tensor)
-    names: tuple[Optional[str], ...] = tensor.names
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    stripped: torch.Tensor = without_dim_names(tensor)
 
     num_chunks: int = cp_size * 2
     chunks: tuple[torch.Tensor, ...] = stripped.chunk(num_chunks, dim=dim)
@@ -97,5 +99,5 @@ def _reorder_zigzag_to_natural(
     result: torch.Tensor = torch.cat([chunks[i] for i in order], dim=dim)
 
     if names[0] is not None:
-        result = result.refine_names(*names)
+        result = apply_dim_names(result, list(names))
     return result

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/concat_steps/executor.py

@@ -7,6 +7,7 @@ import torch
 from sglang.srt.debug_utils.comparator.dims_spec import (
     SEQ_DIM_NAME,
     TOKEN_DIM_NAME,
+    get_dim_names,
 )
 from sglang.srt.debug_utils.comparator.utils import Pair
 
@@ -30,10 +31,9 @@ def execute_token_aligner_concat_steps(
 def _resolve_token_dim(tensor: torch.Tensor) -> int:
     """Find the token/seq dim index. Falls back to dim 0 for unnamed tensors or
     tensors without a recognised token/seq dim."""
-    if tensor.names[0] is None:
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    if names[0] is None:
         return _UNNAMED_TOKEN_DIM_FALLBACK
-
-    names: tuple[Optional[str], ...] = tensor.names
     for candidate in (TOKEN_DIM_NAME, SEQ_DIM_NAME):
         if candidate in names:
             return list(names).index(candidate)

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/aux_loader.py

@@ -29,6 +29,7 @@ from sglang.srt.debug_utils.comparator.dims_spec import (
     TokenLayout,
     apply_dim_names,
     resolve_dim_names,
+    without_dim_names,
 )
 from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank
 from sglang.srt.debug_utils.comparator.log_sink import log_sink
@@ -242,8 +243,8 @@ def _load_and_align_aux_tensor(
 
         sub_result = execute_sub_plans(tensors=tensors, plans=sub_plans)
         assert sub_result.tensor is not None
-        return sub_result.tensor.rename(
-            None
+        return without_dim_names(
+            sub_result.tensor
         )  # strip named dims before returning to plugin
 
     log_sink.add(

python/sglang/srt/debug_utils/comparator/aligner/token_aligner/smart/executor.py

@@ -12,8 +12,9 @@ from sglang.srt.debug_utils.comparator.dims_spec import (
     SEQ_DIM_NAME,
     TOKEN_DIM_NAME,
     TokenLayout,
+    apply_dim_names,
+    get_dim_names,
     resolve_dim_by_name,
-    strip_dim_names,
 )
 from sglang.srt.debug_utils.comparator.utils import Pair
 
@@ -77,16 +78,16 @@ def _collapse_bs_to_t(
         )
 
     lhs_pattern, rhs_pattern, new_names = _build_bs_collapse_pattern(
-        names=list(some_tensor.names),
+        names=list(get_dim_names(some_tensor)),
         batch_dim=batch_dim,
         seq_dim=seq_dim,
     )
 
     result: dict[int, torch.Tensor] = {}
     for step, tensor in tensor_of_step.items():
-        plain: torch.Tensor = strip_dim_names(tensor)
-        collapsed: torch.Tensor = rearrange(plain, f"{lhs_pattern} -> {rhs_pattern}")
-        result[step] = collapsed.refine_names(*new_names)
+        collapsed: torch.Tensor = rearrange(tensor, f"{lhs_pattern} -> {rhs_pattern}")
+        collapsed = apply_dim_names(collapsed, [n for n in new_names if n is not None])
+        result[step] = collapsed
 
     return result
 
@@ -121,7 +122,7 @@ def _build_bs_collapse_pattern(
 
 
 def _resolve_dim_or_fallback(tensor: torch.Tensor, name: str) -> int:
-    if tensor.names[0] is None:
+    if get_dim_names(tensor)[0] is None:
         return _UNNAMED_TOKEN_DIM_FALLBACK
     return resolve_dim_by_name(tensor, name)
 
@@ -143,7 +144,7 @@ def _extract_and_stack_tokens(
     token_dim: int = _resolve_dim_or_fallback(some_tensor, TOKEN_DIM_NAME)
 
     tokens: list[torch.Tensor] = [
-        strip_dim_names(tensor_of_step[s]).select(dim=token_dim, index=i)
+        tensor_of_step[s].select(dim=token_dim, index=i)
         for s, i in zip(locator.steps, locator.token_index_in_step)
     ]
     return torch.stack(tokens, dim=token_dim)

python/sglang/srt/debug_utils/comparator/aligner/unsharder/executor.py

@@ -13,7 +13,10 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
 )
 from sglang.srt.debug_utils.comparator.dims_spec import (
     ParallelAxis,
+    apply_dim_names,
+    get_dim_names,
     resolve_dim_by_name,
+    without_dim_names,
 )
 from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
 from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import compute_diff
@@ -65,7 +68,11 @@ def _apply_unshard(
 
     if isinstance(params, ConcatParams):
         dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
-        return torch.cat(ordered_tensors, dim=dim), []
+        names: tuple[Optional[str], ...] = get_dim_names(ordered_tensors[0])
+        result = torch.cat(ordered_tensors, dim=dim)
+        if names[0] is not None:
+            result = apply_dim_names(result, list(names))
+        return result, []
 
     if isinstance(params, CpThdConcatParams):
         thd_dim: int = resolve_dim_by_name(ordered_tensors[0], params.dim_name)
@@ -79,11 +86,11 @@ def _apply_unshard(
         )
 
     if isinstance(params, ReduceSumParams):
-        stripped: list[torch.Tensor] = [t.rename(None) for t in ordered_tensors]
+        names: tuple[Optional[str], ...] = get_dim_names(ordered_tensors[0])
+        stripped: list[torch.Tensor] = [without_dim_names(t) for t in ordered_tensors]
         result: torch.Tensor = torch.stack(stripped).sum(dim=0)
-        names: tuple[Optional[str], ...] = ordered_tensors[0].names
         if names[0] is not None:
-            result = result.refine_names(*names)
+            result = apply_dim_names(result, list(names))
         return result, []
 
     raise ValueError(f"Unsupported unshard operation: {type(params).__name__}")
@@ -95,7 +102,7 @@ def _verify_replicated_group(
     axis: ParallelAxis,
     group_index: int,
 ) -> list[ReplicatedCheckResult]:
-    baseline: torch.Tensor = ordered_tensors[0].rename(None).float()
+    baseline: torch.Tensor = ordered_tensors[0].float()
 
     return [
         _check_replicated_pair(
@@ -117,7 +124,7 @@ def _check_replicated_pair(
     group_index: int,
     compared_index: int,
 ) -> ReplicatedCheckResult:
-    other_float: torch.Tensor = other.rename(None).float()
+    other_float: torch.Tensor = other.float()
 
     if baseline.shape != other_float.shape:
         passed = False
@@ -155,8 +162,8 @@ def _thd_concat(
     This function splits each rank by seq_lens, then interleaves across ranks
     per-seq: [seqA_r0 + seqA_r1 + ... | seqB_r0 + seqB_r1 + ... | tail_pad].
     """
-    names: tuple[Optional[str], ...] = ordered_tensors[0].names
-    stripped: list[torch.Tensor] = [t.rename(None) for t in ordered_tensors]
+    names: tuple[Optional[str], ...] = get_dim_names(ordered_tensors[0])
+    stripped: list[torch.Tensor] = [without_dim_names(t) for t in ordered_tensors]
 
     # Split each rank into [seq0, seq1, ..., tail_remainder]
     split_sizes: list[int] = list(seq_lens_per_rank)
@@ -179,5 +186,5 @@ def _thd_concat(
     )
 
     if names[0] is not None:
-        result = result.refine_names(*names)
+        result = apply_dim_names(result, list(names))
     return result

python/sglang/srt/debug_utils/comparator/bundle_comparator.py

@@ -23,8 +23,10 @@ from sglang.srt.debug_utils.comparator.dims_spec import (
     TOKEN_DIM_NAME,
     ParallelAxis,
     apply_dim_names,
+    get_dim_names,
     parse_dims,
     resolve_dim_names,
+    without_dim_names,
 )
 from sglang.srt.debug_utils.comparator.dp_utils import filter_to_non_empty_dp_rank
 from sglang.srt.debug_utils.comparator.log_sink import log_sink
@@ -298,8 +300,8 @@ def _compare_bundle_pair_tensor_type(
     )
 
     # Compare
-    aligned_baseline: torch.Tensor = aligner_result.tensors.x.rename(None)
-    aligned_target: torch.Tensor = aligner_result.tensors.y.rename(None)
+    aligned_baseline: torch.Tensor = without_dim_names(aligner_result.tensors.x)
+    aligned_target: torch.Tensor = without_dim_names(aligner_result.tensors.y)
 
     info = compare_tensor_pair(
         x_baseline=aligned_baseline,
@@ -361,10 +363,9 @@ def _try_generate_viz(
 
 def _resolve_seq_dim(tensor: torch.Tensor) -> Optional[int]:
     """Find the token/seq dimension index from the tensor's named dims."""
-    if tensor.names[0] is None:
+    names: tuple[Optional[str], ...] = get_dim_names(tensor)
+    if names[0] is None:
         return None
-
-    names: tuple[Optional[str], ...] = tensor.names
     for target_name in (TOKEN_DIM_NAME, SEQ_DIM_NAME):
         if target_name in names:
             return list(names).index(target_name)

python/sglang/srt/debug_utils/comparator/dims_spec/__init__.py

@@ -7,8 +7,9 @@ from sglang.srt.debug_utils.comparator.dims_spec.dims_parser import (
 from sglang.srt.debug_utils.comparator.dims_spec.tensor_naming import (
     apply_dim_names,
     find_dim_index,
+    get_dim_names,
     resolve_dim_by_name,
-    strip_dim_names,
+    without_dim_names,
 )
 from sglang.srt.debug_utils.comparator.dims_spec.types import (
     _FUSED_NAME_SEP,
@@ -41,9 +42,10 @@ __all__ = [
     "_SingletonDimUtil",
     "apply_dim_names",
     "find_dim_index",
+    "get_dim_names",
     "parse_dim",
     "parse_dims",
     "resolve_dim_by_name",
     "resolve_dim_names",
-    "strip_dim_names",
+    "without_dim_names",
 ]

python/sglang/srt/debug_utils/comparator/dims_spec/tensor_naming.py

@@ -6,6 +6,8 @@ import torch
 
 from sglang.srt.debug_utils.comparator.dims_spec.types import DimSpec
 
+_DIM_NAMES_ATTR = "_dim_names"
+
 
 def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:
     """Find index by name. Accepts both ``*``-form and ``___``-form for fused dims."""
@@ -15,11 +17,22 @@ def find_dim_index(dim_specs: list[DimSpec], name: str) -> Optional[int]:
     return None
 
 
+def get_dim_names(tensor: torch.Tensor) -> tuple[Optional[str], ...]:
+    """Get dimension names attached to a tensor.
+
+    Returns a tuple of ``None`` values if no names are attached.
+    """
+    names = getattr(tensor, _DIM_NAMES_ATTR, None)
+    if names is not None:
+        return names
+    return (None,) * tensor.ndim
+
+
 def resolve_dim_by_name(tensor: torch.Tensor, name: str) -> int:
-    if tensor.names[0] is None:
+    names = get_dim_names(tensor)
+    if names[0] is None:
         raise ValueError(f"Tensor has no names, cannot resolve {name!r}")
 
-    names: tuple[Optional[str], ...] = tensor.names
     try:
         return list(names).index(name)
     except ValueError:
@@ -33,8 +46,11 @@ def apply_dim_names(tensor: torch.Tensor, dim_names: list[str]) -> torch.Tensor:
             f"but dims string specifies {len(dim_names)} names {dim_names}. "
             f"Please fix the dims string in the dumper.dump() call to match the actual tensor shape."
         )
-    return tensor.refine_names(*dim_names)
+    view = torch.ops.aten.alias(tensor)
+    view._dim_names = tuple(dim_names)
+    return view
 
 
-def strip_dim_names(tensor: torch.Tensor) -> torch.Tensor:
-    return tensor.rename(None)
+def without_dim_names(tensor: torch.Tensor) -> torch.Tensor:
+    # Returns a new view without _dim_names; the original tensor is not modified.
+    return torch.ops.aten.alias(tensor)