📅 创建时间：2026-06-03 🏷️ 标签：#动态Shape #SymbolicShape #PartialShape #DynamicBatch #Broadcast #PolymorphicShapes 📚 前置知识：[[/04-ai/01-llm-engineering/07-llm-evolution]]（LLM 发展脉络） 📚 前置知识：[[07-graph-optimization-passes]]（图优化 Pass） 📚 前置知识：[[09-memory-planning]]（内存规划） 📚 相关知识：[[03-graph-representation]]（计算图表示） 📚 相关知识：[[24-torch-compile]]（torch.compile）

┌──────────────────────────────────────────────────────────────────────────────┐ │ 🔍 场景：torch.compile 报错 "Unable to specialize type of tensor" │ ├──────────────────────────────────────────────────────────────────────────────┤ │ 你在部署 LLM 推理服务，需要支持任意长度的 prompt： │ │ - 短 prompt：32 tokens，延迟 10ms │ │ - 长 prompt：2048 tokens，延迟 150ms │ │ - 极长 prompt：8192 tokens，OOM 了 │ │ │ │ 你尝试用 torch.compile 编译一个可变长度的 transformer，但编译器报错： │ │ "Unable to specialize type of tensor with dynamic dimension. │ │ Hint: the dimension with value None may need bounds." │ │ │ │ 编译器怎么理解"不确定的 shape"？为什么动态 Shape 这么难处理？ │ └──────────────────────────────────────────────────────────────────────────────┘

动态 Shape——符号分析与形状处理 / Dynamic Shapes, Symbolic Analysis, and Shape Processing

第1节：动态 Shape 的挑战

1.1 为什么动态 Shape 困难

python

# 静态编译 vs 动态 Shape 的对比

# 静态编译（已知 shape）：
# 输入: (batch=4, seq_len=512, hidden=768)
# 
# 编译器可以：
# 1. 预分配固定大小的显存
# 2. 展开所有循环
# 3. 生成高效的机器码
# 4. 做各种静态优化

# 动态 Shape（未知 shape）：
# 输入: (batch=B, seq_len=L, hidden=768)
# 其中 B, L 是运行时才知道的值
#
# 编译器无法：
# 1. 预分配显存（不知道需要多大）
# 2. 展开循环（不知道迭代次数）
# 3. 生成高效的机器码（需要根据 shape 特殊化）
# 4. 做很多静态优化

1.2 动态 Shape 的典型场景

python

# 场景1：LLM 推理
class DynamicLLMInference:
    """
    LLM 推理中的动态 Shape
    
    问题：
    1. prompt 长度不确定（32 ~ 8192）
    2. 生成 token 数不确定（1 ~ 2048）
    3. KV Cache 大小随长度线性增长
    """
    
    def forward(self, input_ids, max_length):
        """
        input_ids: (batch_size, prompt_length)
        运行时 prompt_length 才知道
        """
        pass


# 场景2：Dynamic Batch
class DynamicBatchInference:
    """
    动态 Batch 推理
    
    问题：
    1. 每个请求的 batch size 不同
    2. 需要 padding 到统一长度
    3. 效率损失
    """
    
    def forward(self, batched_inputs):
        """
        batched_inputs: list of tensors, each (seq_len, hidden)
        实际运行时 batch size 才知道
        """
        pass


# 场景3：变长序列处理
class VariableLengthProcessing:
    """
    变长序列处理
    
    典型问题：
    1. NLP 中的句子长度不同
    2. 语音处理中的音频长度不同
    3. 时序数据的窗口大小不同
    """
    
    def forward(self, sequences):
        """
        sequences: 每个序列长度不同
        无法用固定 shape 表示
        """
        pass

1.3 动态 Shape 的类型

python

from typing import Optional, Union

class ShapeTypes:
    """
    Shape 的几种类型
    
    1. 静态 Shape：编译期完全确定
    2. 动态 Shape（无界）：只知道部分信息
    3. 动态 Shape（有界）：有上下界约束
    4. Polymorphic Shape：符号化的 Shape
    """
    
    # 静态 Shape
    static_shape = (4, 512, 768)  # 完全确定
    
    # 动态 Shape（无界）- None 表示动态
    dynamic_unbounded = (4, None, 768)  # seq_len 未知
    
    # 动态 Shape（有界）- 使用 Range
    dynamic_bounded = (4, Range(1, 8192), 768)  # seq_len 在 1~8192 之间
    
    # 符号 Shape
    symbolic_shape = (Batch, SeqLen, 768)  # Batch, SeqLen 是符号
    
    # 部分 Shape
    partial_shape = (4, Dimension(divisible_by=64), 768)  # 知道能被 64 整除


class Range:
    """带约束的维度范围"""
    
    def __init__(self, lower: int, upper: int):
        self.lower = lower
        self.upper = upper
    
    def contains(self, value: int) -> bool:
        return self.lower <= value <= self.upper

第2节：Symbolic Shape——符号化维度

2.1 符号维度的表示

python

class Symbol:
    """
    符号维度
    
    用途：
    1. 表示编译期未知的维度
    2. 跟踪维度之间的关系
    3. 支持符号计算
    """
    
    _counter = 0
    
    def __init__(self, name: str = None, constraints: dict = None):
        self.id = Symbol._counter
        Symbol._counter += 1
        self.name = name or f"s{self.id}"
        self.constraints = constraints or {}
    
    def __repr__(self):
        return f"${self.name}"
    
    def __hash__(self):
        return hash(self.id)
    
    def __eq__(self, other):
        if isinstance(other, Symbol):
            return self.id == other.id
        return False


class SymbolicDim:
    """
    符号化维度
    
    支持的操作：
    1. 整数运算：+, -, *, //, %
    2. 比较：==, !=, <, <=, >, >=
    3. 约束传播
    """
    
    def __init__(self, value: Union[int, Symbol, 'SymbolicDim']):
        if isinstance(value, SymbolicDim):
            self.value = value.value
            self.constraints = value.constraints.copy()
        elif isinstance(value, Symbol):
            self.value = value
            self.constraints = {}
        else:
            self.value = int(value)
            self.constraints = {}
    
    def __add__(self, other):
        result = SymbolicDim(self.value)
        if isinstance(other, SymbolicDim):
            if isinstance(self.value, int) and isinstance(other.value, int):
                result.value = self.value + other.value
            else:
                result.value = Symbol(f"add_{self}_{other}")
        else:
            result.value = self.value + other
        return result
    
    def __mul__(self, other):
        result = SymbolicDim(self.value)
        if isinstance(other, SymbolicDim):
            if isinstance(self.value, int) and isinstance(other.value, int):
                result.value = self.value * other.value
        return result
    
    def __floordiv__(self, other):
        result = SymbolicDim(self.value)
        if isinstance(self.value, int) and isinstance(other, int):
            result.value = self.value // other
        return result


# 示例
batch = SymbolicDim(Symbol('batch'))
seq_len = SymbolicDim(Symbol('seq_len'))
hidden = SymbolicDim(768)

# 符号表达式
output_seq = seq_len // 2  # 下采样后的序列长度
attn_scores = batch * seq_len * seq_len  # attention score 的大小

2.2 约束传播

python

class ConstraintPropagation:
    """
    约束传播
    
    核心思想：
    1. 从输入约束推导中间值的约束
    2. 从中间值约束推导输出约束
    3. 传播可以揭示更多关系或检测矛盾
    """
    
    def __init__(self):
        self.constraints = {}  # symbol -> [constraints]
    
    def add_constraint(self, sym: Symbol, constraint):
        """添加约束"""
        if sym not in self.constraints:
            self.constraints[sym] = []
        self.constraints[sym].append(constraint)
    
    def propagate_conv(self, input_dim, kernel_size, stride, padding):
        """
        Conv2D 的约束传播
        
        output = (input - kernel + 2*padding) // stride + 1
        """
        # 输入维度约束
        # 已知 output, stride, padding, kernel，推导 input
        min_input = (output - 1) * stride - 2 * padding + kernel
        
        # 约束传播
        return {
            'min': min_input,
            'divisible_by': stride,  # input - kernel + 2*padding ≡ 0 (mod stride)
        }
    
    def propagate_matmul(self, a_dim, b_dim):
        """
        MatMul 的约束传播
        
        (M, K) @ (K, N) = (M, N)
        
        中间维度 K 必须相等
        """
        if isinstance(a_dim, SymbolicDim) and isinstance(b_dim, SymbolicDim):
            # K 必须相等
            self.unify(a_dim, b_dim)
    
    def unify(self, dim1, dim2):
        """统一两个维度"""
        # 如果都是符号，要求它们相等
        # 如果一个是具体值，设置另一个的约束
        pass

2.3 Symbolic Shape 的实际应用

python

# JAX/XLA 风格的 Polymorphic Shapes
class PolymorphicShapesExample:
    """
    JAX 的 shape polymorphism 支持
    
    使用符号维度表示动态 shape
    """
    
    def jax_style_function(self, x, y):
        """
        JAX 支持的动态 shape 函数签名
        
        def f(x: float32[B, H], y: float32[H, D]) -> float32[B, D]:
            return x @ y
        
        其中 B, H, D 是符号维度
        """
        pass
    
    def torch_style_with_constraints(self, x, y):
        """
        torch.compile 的 dynamic=True 模式
        
        @torch.compile(dynamic=True)
        def f(x, y):
            return x @ y
        
        编译器会：
        1. 用符号维度 B, H, D 替换具体值
        2. 生成可以在任意 B, H, D 下运行的代码
        3. 保留 shape specialization 的能力
        """
        pass


# Symbolic Shape 调试
def debug_symbolic_shape():
    """调试符号维度"""
    
    import torch
    
    # 定义动态 shape
    x = torch.randn(4, None, 768)  # None 表示动态
    print(f"Shape: {x.shape}")
    print(f"Dims: {x.dim()}")
    
    # 编译时检查
    compiled = torch.compile(lambda x: x @ torch.randn(768, 10), dynamic=True)
    
    # 测试不同 shape
    for seq_len in [32, 64, 128, 512]:
        x = torch.randn(4, seq_len, 768)
        out = compiled(x)
        print(f"seq_len={seq_len}, output shape: {out.shape}")

第3节：Partial Shape——部分已知形状

3.1 Partial Shape 的表示

python

from dataclasses import dataclass
from typing import Optional, List

@dataclass
class PartialShape:
    """
    部分 Shape
    
    表示只知道部分维度的 shape
    """
    
    dims: List[Optional[int]]  # None 表示未知
    
    def __repr__(self):
        return f"PartialShape({self.dims})"
    
    @staticmethod
    def fully_static(dims: List[int]) -> 'PartialShape':
        return PartialShape(dims)
    
    @staticmethod
    def fully_dynamic(rank: int) -> 'PartialShape':
        return PartialShape([None] * rank)
    
    @staticmethod
    def from_onnx(onnx_tensor_type):
        """从 ONNX tensor type 转换"""
        shape = onnx_tensor_type.shape
        dims = []
        for dim in shape.dim:
            if dim.HasField('dim_param'):
                dims.append(None)  # 符号维度
            elif dim.HasField('dim_value'):
                dims.append(dim.dim_value)  # 具体值
            else:
                dims.append(None)  # 完全未知
        return PartialShape(dims)
    
    def is_fully_static(self) -> bool:
        return all(d is not None for d in self.dims)
    
    def is_fully_dynamic(self) -> bool:
        return all(d is None for d in self.dims)
    
    def get_static_dims(self) -> List[int]:
        return [d for d in self.dims if d is not None]
    
    def get_dynamic_dim_indices(self) -> List[int]:
        return [i for i, d in enumerate(self.dims) if d is None]


# Partial Shape 示例
examples = [
    # 静态
    PartialShape([4, 512, 768]),
    
    # 只有 batch 动态
    PartialShape([None, 512, 768]),
    
    # batch 和 seq_len 动态
    PartialShape([None, None, 768]),
    
    # 只有 hidden 静态
    PartialShape([None, None, 768]),
]

3.2 Partial Shape 的约束

python

@dataclass
class DynamicDimConstraint:
    """
    动态维度的约束
    """
    
    # 下界
    min_value: Optional[int] = None
    
    # 上界
    max_value: Optional[int] = None
    
    # 必须整除的值
    divisible_by: Optional[int] = None
    
    # 必须等于另一个维度
    equal_to: Optional[int] = None
    
    # 必须是的倍数
    multiple_of: Optional[int] = None
    
    # 凤巢约束
    def implies(self, other: 'DynamicDimConstraint') -> bool:
        """检查 self 是否蕴含 other"""
        if other.min_value and self.min_value > other.min_value:
            return False
        if other.max_value and self.max_value < other.max_value:
            return False
        return True
    
    def merge(self, other: 'DynamicDimConstraint') -> 'DynamicDimConstraint':
        """合并两个约束"""
        return DynamicDimConstraint(
            min_value=max(self.min_value, other.min_value) if self.min_value and other.min_value else (self.min_value or other.min_value),
            max_value=min(self.max_value, other.max_value) if self.max_value and other.max_value else (self.max_value or other.max_value),
            divisible_by=self._lcm(self.divisible_by, other.divisible_by),
            multiple_of=self._gcd(self.multiple_of, other.multiple_of),
        )
    
    def _lcm(self, a, b):
        if a is None: return b
        if b is None: return a
        return a * b // self._gcd(a, b)
    
    def _gcd(self, a, b):
        if a is None: return b
        if b is None: return a
        while b: a, b = b, a % b
        return a


class ConstraintDatabase:
    """
    约束数据库
    
    管理所有动态维度的约束
    """
    
    def __init__(self):
        self.constraints = {}  # dim_index -> DynamicDimConstraint
    
    def set_constraint(self, dim_index: int, constraint: DynamicDimConstraint):
        self.constraints[dim_index] = constraint
    
    def add_constraint(self, dim_index: int, constraint: DynamicDimConstraint):
        if dim_index not in self.constraints:
            self.constraints[dim_index] = constraint
        else:
            self.constraints[dim_index] = self.constraints[dim_index].merge(constraint)
    
    def check_feasibility(self) -> bool:
        """检查约束是否可满足"""
        for dim_idx, constraint in self.constraints.items():
            if constraint.min_value and constraint.max_value:
                if constraint.min_value > constraint.max_value:
                    return False  # 矛盾
        return True
    
    def get_constraint(self, dim_index: int) -> DynamicDimConstraint:
        return self.constraints.get(dim_index, DynamicDimConstraint())

第4节：Shape Specialization——形状特化

4.1 编译期特化

python

class ShapeSpecialization:
    """
    Shape 特化
    
    策略：
    1. 编译期选择特化某些维度
    2. 其他维度保持动态
    3. 权衡特化收益和代码膨胀
    """
    
    def specialize(self, graph, strategy):
        """
        对图进行 shape 特化
        
        策略选项：
        - 'full_static': 完全静态
        - 'batch_dynamic': batch 动态，其他静态
        - 'seq_dynamic': seq_len 动态，其他静态
        - 'full_dynamic': 完全动态
        """
        
        if strategy == 'batch_dynamic':
            return self._specialize_batch_dynamic(graph)
        elif strategy == 'seq_dynamic':
            return self._specialize_seq_dynamic(graph)
        elif strategy == 'full_dynamic':
            return self._specialize_full_dynamic(graph)
    
    def _specialize_batch_dynamic(self, graph):
        """batch 动态，其他静态"""
        for node in graph.nodes:
            for i, dim in enumerate(node.output_shape.dims):
                if i == 0:  # batch 维度
                    node.output_shape.dims[i] = None  # 动态
                else:
                    # 尝试推导具体值
                    inferred = self._infer_dim(node, i)
                    node.output_shape.dims[i] = inferred
        return graph
    
    def _infer_dim(self, node, dim_index):
        """
        推断特定维度的具体值
        """
        if node.op_type == 'matmul':
            # matmul: output[...,-1] = input2.shape[-1]
            if dim_index == len(node.output_shape.dims) - 1:
                return node.inputs[1].shape[-1]
        elif node.op_type == 'reshape':
            # reshape: 输出维度的乘积 = 输入维度的乘积
            # 某些维度可能可以推断
            pass
        return None  # 无法推断，保持动态

4.2 JIT 特化

python

# 运行时特化（JIT Compilation）

class JITSpecialization:
    """
    JIT 形状特化
    
    运行时根据实际 shape 生成 specialized code
    """
    
    def __init__(self):
        self.cache = {}  # (shape_tuple) -> compiled_code
    
    def compile(self, fn, shape):
        """
        根据 shape 编译函数
        """
        shape_key = tuple(shape)
        
        if shape_key in self.cache:
            return self.cache[shape_key]
        
        # 生成 specialized code
        specialized = self._specialize(fn, shape)
        
        # 编译
        compiled = torch.compile(specialized, mode='reduce-overhead')
        
        self.cache[shape_key] = compiled
        return compiled
    
    def _specialize(self, fn, shape):
        """
        生成特化版本的函数
        """
        # 可以在函数中嵌入 shape 信息
        # 帮助编译器做更多优化
        return fn

第5节：Dynamic Dimension Bounds——动态维度边界

5.1 为什么要 Bounds

python

# 为什么需要给动态维度加边界

# 问题：没有边界时
# def process(x):
#     # x.shape = (B, L, 768)
#     # B 和 L 都是动态的
#     
#     # 编译器无法：
#     # 1. 检查 B*L*768*4 > GPU_memory 是否成立
#     # 2. 确定是否需要 OOM 检查
#     # 3. 优化内存分配
#     
#     y = x @ weight  # 无法确定中间 tensor 大小
#     return y

# 解决方案：提供边界
# @torch.compile(dynamic=True)
# def process(x):
#     assert x.shape[1] >= 1 and x.shape[1] <= 8192  # seq_len 边界
#     assert x.shape[0] >= 1 and x.shape[0] <= 256   # batch 边界
#     
#     # 编译器现在知道：
#     # 1. 最大显存需求
#     # 2. 可以做 OOM 检查
#     # 3. 可以优化内存分配
#     
#     y = x @ weight
#     return y

5.2 Bounds 传播

python

class BoundsPropagation:
    """
    边界传播
    
    从输入边界推导所有中间值的边界
    """
    
    def propagate(self, graph, input_bounds):
        """
        传播边界
        
        input_bounds: {
            0: DynamicDimConstraint(min=1, max=256),    # batch
            1: DynamicDimConstraint(min=1, max=8192),  # seq_len
            2: DynamicDimConstraint(min=768, max=768), # hidden (static)
        }
        """
        
        bounds = {}
        bounds.update(input_bounds)
        
        for node in topological_sort(graph):
            output_bounds = self._compute_node_bounds(node, bounds)
            bounds[node.id] = output_bounds
        
        return bounds
    
    def _compute_node_bounds(self, node, input_bounds):
        """计算节点的输出边界"""
        
        if node.op_type == 'matmul':
            # output_shape = (B, K2) where:
            # B = input1.shape[:-1]
            # K2 = input2.shape[-1]
            
            # 假设 input1 = (B, K), input2 = (K, K2)
            input0_shape = input_bounds.get(node.inputs[0].id, [])
            input1_shape = input_bounds.get(node.inputs[1].id, [])
            
            # 推导输出边界
            # 第一个维度和输入第一个维度相同（通常是 batch）
            # 最后一个维度是第二个输入的最后一个维度
            return input0_shape[:-1] + [input1_shape[-1]]
        
        elif node.op_type == 'softmax':
            # softmax 不改变 shape
            return input_bounds.get(node.inputs[0].id, [])
        
        elif node.op_type == 'reshape':
            # reshape: 输出维度的乘积 = 输入维度的乘积
            input_shape = input_bounds.get(node.inputs[0].id, [])
            target_shape = node.attrs['shape']
            
            # 推导新边界
            return self._propagate_reshape_bounds(input_shape, target_shape)
        
        return [DynamicDimConstraint()] * node.output_rank
    
    def _propagate_reshape_bounds(self, input_shape, target_shape):
        """推导 reshape 后的边界"""
        
        output_bounds = []
        
        for dim in target_shape:
            if isinstance(dim, int):
                output_bounds.append(DynamicDimConstraint(
                    min_value=dim, max_value=dim
                ))
            elif dim == -1:
                # 需要从其他维度推导
                known_product = 1
                unknown_count = 0
                for i, d in enumerate(target_shape):
                    if isinstance(d, int):
                        known_product *= d
                    elif i != target_shape.index(-1):
                        if input_shape[i]:
                            known_product *= input_shape[i].max_value or 1
                
                output_bounds.append(DynamicDimConstraint(
                    min_value=1, max_value=known_product
                ))
            else:
                output_bounds.append(DynamicDimConstraint())
        
        return output_bounds

5.3 Bounds 的使用场景

python

# Bounds 的实际使用

class BoundsUsage:
    """
    如何使用 bounds 进行优化
    """
    
    @staticmethod
    def memory_estimation(bounds):
        """
        使用 bounds 估算显存需求
        
        bounds: 每个维度的约束
        
        估算最大显存需求：
        max_memory = Π(max_dim for dim in bounds) * element_size
        """
        
        max_elements = 1
        for bound in bounds:
            if bound.max_value:
                max_elements *= bound.max_value
            else:
                # 没有上界，无法估算
                return None
        
        return max_elements * 4  # float32 = 4 bytes
    
    @staticmethod
    def check_oom_possible(bounds, available_memory):
        """
        检查是否可能 OOM
        """
        max_memory = BoundsUsage.memory_estimation(bounds)
        if max_memory is None:
            return True  # 无法确定，可能 OOM
        return max_memory > available_memory
    
    @staticmethod
    def optimize_allocation(bounds):
        """
        使用 bounds 优化显存分配
        
        知道最大需求后，可以预分配
        """
        max_elements = 1
        for bound in bounds:
            if bound.max_value:
                max_elements *= bound.max_value
        
        # 预分配
        buffer = torch.empty(max_elements, device='cuda')
        return buffer

第6节：动态 Shape 下的优化限制

6.1 无法做的优化

python

class DynamicShapeOptimizationLimitations:
    """
    动态 Shape 下无法进行的优化
    """
    
    OPTIMIZATIONS_NOT_AVAILABLE = {
        'StaticMemoryAllocation': """
            无法在编译期预分配固定大小的显存。
            必须使用动态分配或 arena 分配器。
        """,
        
        'FullLoopUnrolling': """
            循环次数在编译期未知，无法完全展开。
            只能做部分展开（最多展开的次数 = 最大迭代次数）。
        """,
        
        'CompleteKernelFusion': """
            某些 fusion 需要知道 tensor 大小。
            如 reduce 操作后的 reshape。
        """,
        
        'MemoryLayoutOptimization': """
            无法确定最优的 tensor layout。
            只能在运行时选择或保守选择通用 layout。
        """,
        
        'DeadCodeEliminationForDynamicPaths': """
            如果 shape 决定控制流，编译期无法确定哪些路径会被执行。
            无法删除永远不会执行的代码。
        """,
    }
    
    @staticmethod
    def get_available_optimizations():
        """
        动态 Shape 下仍然可用的优化
        """
        return {
            'PartialFusion': "部分算子可以融合",
            'LayoutSelection': "可以根据实际 shape 选择 layout",
            'LazyAllocation': "延迟分配，只在需要时分配",
            'StreamingMemory': "流式处理，分块计算",
        }

6.2 权衡策略

python

class DynamicShapeTradeoffs:
    """
    动态 Shape 的权衡
    """
    
    STRATEGIES = {
        'full_static': {
            'description': '选择最大 shape 进行静态编译',
            'pros': ['最优性能', '完整优化'],
            'cons': ['显存浪费', '不支持更小 shape'],
            'use_case': '生产环境，shape 固定'
        },
        
        'full_dynamic': {
            'description': '保持完全动态',
            'pros': ['灵活', '支持任意 shape'],
            'cons': ['性能损失 20-50%', '无法做很多优化'],
            'use_case': '开发/调试，或 shape 高度变化'
        },
        
        'bounded_dynamic': {
            'description': '提供边界约束，保持动态',
            'pros': ['灵活 + 部分优化', '可以做 OOM 检查'],
            'cons': ['需要用户提供边界', '仍有部分性能损失'],
            'use_case': '生产环境，shape 有范围'
        },
        
        'multi_version': {
            'description': '编译多个版本，每个版本对应一个 shape',
            'pros': ['每个版本最优', '仍支持动态'],
            'cons': ['编译时间长', '代码膨胀'],
            'use_case': 'shape 只有几种固定值'
        },
    }
    
    def choose_strategy(self, shape_variation, performance_requirement):
        """
        选择合适的策略
        """
        if len(shape_variation) <= 4:
            return 'multi_version'
        elif all(s.max_value - s.min_value < 10 for s in shape_variation.values()):
            return 'bounded_dynamic'
        elif performance_requirement == 'maximum':
            return 'full_static'
        else:
            return 'full_dynamic'

第7节：JAX/XLA 的动态 Shape 支持

7.1 JAX 的 Polymorphic Shapes

python

# JAX/XLA 的动态 Shape 支持

# JAX 支持用符号表示 shape
from jax import numpy as jnp

def jax_polymorphic_function(x, y):
    """
    JAX 的 polymorphic shapes
    
    def f(x: float32[B, H], y: float32[H, D]) -> float32[B, D]:
        return x @ y
    
    B, H, D 是符号维度，可以在运行时是任意值
    """
    return x @ y


# 使用 shape polymorphism
from jax.experimental import export

# 导出的 API 允许指定 shape 范围
exported = export.export(
    jax_polymorphic_function,
    polymorphic_shapes=['(batch, 768)', '(768, 10)']
)

# batch 是动态的，可以是任意正整数

7.2 Triton 的动态 Shape 支持

python

import triton

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64}),
        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 64}),
        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}),
    ],
    key=['M', 'N'],  # 根据 M, N 选择配置
)
@triton.jit
def triton_dynamic_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    """
    Triton 支持动态 shape
    
    M, N, K 是运行时才知道的值
    Triton 会根据实际值选择最优的配置
    """
    pass


def call_triton_dynamic(x, y):
    """
    调用动态 Triton kernel
    """
    M, K = x.shape
    K2, N = y.shape
    assert K == K2
    
    c = torch.empty((M, N), device=x.device, dtype=x.dtype)
    
    # M, N 是动态的，Triton 运行时选择最优配置
    grid = lambda META: (triton.cdiv(M, META['BLOCK_M']), 
                         triton.cdiv(N, META['BLOCK_N']))
    
    triton_dynamic_kernel[grid](
        x, y, c, M, N, K, 
        x.stride(0), x.stride(1),
        y.stride(0), y.stride(1),
        c.stride(0), c.stride(1),
    )
    
    return c

第8节：PyTorch 2.0 的动态 Shape 处理

8.1 torch.compile 的 dynamic=True

python

import torch

# torch.compile 的动态 shape 支持

model = TransformerModel()

# 方式1：完全动态
compiled_model = torch.compile(model, dynamic=True)

# 方式2：指定哪些维度是动态的
compiled_model = torch.compile(
    model, 
    dynamic=True,
    options={'shape_padding': True}
)

# 方式3：调试动态 shape 问题
compiled_model = torch.compile(
    model,
    dynamic=True,
    backend='aot_eager',  # 使用 eager 模式调试
)

# 测试
for batch_size in [1, 4, 16]:
    for seq_len in [32, 128, 512]:
        x = torch.randint(0, 1000, (batch_size, seq_len))
        y = compiled_model(x)
        print(f"batch={batch_size}, seq={seq_len}, out={y.shape}")

8.2 常见错误和解决方案

python

# 动态 Shape 常见错误

# 错误1：维度边界不够
try:
    @torch.compile(dynamic=True)
    def fn(x):
        # 假设 x.shape[1] 会被 512 整除
        # 但编译器需要明确的边界
        return x.view(-1, 512)  # 如果无法被 512 整除会报错
except Exception as e:
    print(f"Error: {e}")
    # 解决方案：添加断言或使用已知的边界

# 解决方案1：使用 assertion
@torch.compile(dynamic=True)
def fn_fixed(x):
    torch._check(x.shape[1] % 512 == 0)  # 断言
    return x.view(-1, 512)

# 解决方案2：提供已知的边界
@torch.compile(dynamic=True)
def fn_with_bounds(x):
    # 编译器知道 seq_len <= 2048
    # 可以做更好的优化
    assert x.shape[1] <= 2048
    return x.view(-1, 512)


# 错误2：shape 相关的 control flow
@torch.compile(dynamic=True)
def fn_control_flow(x):
    if x.shape[0] > 10:  # 运行时才知道
        return x * 2
    else:
        return x * 3
# 这个通常可以工作，但某些情况下可能导致 specialization


# 错误3：不兼容的 shape 操作
@torch.compile(dynamic=True)
def fn_incompatible(x, y):
    # 如果 x.shape[0] 和 y.shape[0] 不相等但都是动态的
    # reshape 可能失败
    return x.reshape(-1) + y.reshape(-1)

8.3 调试动态 Shape 问题

python

# 调试动态 Shape 问题

# 方法1：查看编译日志
torch._dynamo.config.verbose = True
torch._inductor.config.debug = True

model = TransformerModel()
compiled = torch.compile(model, dynamic=True)
output = compiled(input_tensor)

# 方法2：使用 export
from torch.export import export

exported = export(model, args=(input_tensor,), dynamic_shapes={
    'x': {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}
})
print(exported.graph)

# 方法3：捕获错误时的 shape 信息
@torch.compile(dynamic=True, fullgraph=True)
def fn_debug(x):
    print(f"Compiling with x.shape = {x.shape}")  # 打印 shape
    return x @ weight

# 方法4：使用 FX tracer 手动检查
from torch.fx import symbolic_trace

traced = symbolic_trace(model)
print(traced.graph)
for node in traced.graph.nodes:
    if 'shape' in str(node.meta):
        print(f"{node.name}: {node.meta['shape']}")

升华

┌─────────────────────────────────────────────────────────────────────────────┐
│                        动态 Shape 核心原则                                   │
├─────────────────────────────────────────────────────────────────────────────┤
│  1. 动态 Shape 的本质是信息缺失                                              │
│     → 编译器不知道运行时 shape → 无法做很多优化                               │
│  2. 约束是最好的朋友                                                         │
│     → 提供 bounds（上下界）可以解锁很多优化                                   │
│  3. 没有免费的午餐                                                           │
│     → 完全动态 vs 完全静态 需要权衡                                          │
│  4. 多版本编译是性能关键                                                     │
│     → 如果 shape 只有几种值，编译多个版本收益很大                             │
│  5. 边界传播是隐形的优化器                                                  │
│     → 好的边界约束可以让编译器做更多优化                                      │
└─────────────────────────────────────────────────────────────────────────────┘

"AI 可查 vs 必须理解"清单

必须理解（不理解就等于不会）：

🔴 动态 Shape 为什么困难：编译期不知道 shape → 无法预分配显存、无法展开循环、无法做很多优化
🔴 Symbolic Shape 的概念：用符号表示未知维度，支持符号计算和约束传播
🔴 Bounds（边界约束）的价值：提供 min/max 可以解锁 OOM 检查、内存优化、loop unrolling
🔴 Partial Shape：只知道部分维度，其他维度动态
🔴 动态 Shape 下的优化限制：哪些优化做不了，哪些仍可以做

AI 可查（知道去哪查就行）：

✅ JAX/XLA 的 shape polymorphism 具体语法
✅ torch.compile dynamic 的具体配置选项
✅ Triton 动态 shape kernel 的写法
✅ 特定框架的 dynamic shape 最佳实践
✅ 动态 shape 相关的已知 bug 和 workaround

学习状态：🟡 开始学习

动态 Shape——符号分析与形状处理 / Dynamic Shapes, Symbolic Analysis, and Shape Processing ​

第1节：动态 Shape 的挑战 ​

1.1 为什么动态 Shape 困难 ​

1.2 动态 Shape 的典型场景 ​

1.3 动态 Shape 的类型 ​

第2节：Symbolic Shape——符号化维度 ​

2.1 符号维度的表示 ​

2.2 约束传播 ​

2.3 Symbolic Shape 的实际应用 ​

第3节：Partial Shape——部分已知形状 ​

3.1 Partial Shape 的表示 ​

3.2 Partial Shape 的约束 ​

第4节：Shape Specialization——形状特化 ​

4.1 编译期特化 ​

4.2 JIT 特化 ​

第5节：Dynamic Dimension Bounds——动态维度边界 ​

5.1 为什么要 Bounds ​

5.2 Bounds 传播 ​

5.3 Bounds 的使用场景 ​

第6节：动态 Shape 下的优化限制 ​

6.1 无法做的优化 ​

6.2 权衡策略 ​

第7节：JAX/XLA 的动态 Shape 支持 ​

7.1 JAX 的 Polymorphic Shapes ​

7.2 Triton 的动态 Shape 支持 ​

第8节：PyTorch 2.0 的动态 Shape 处理 ​

8.1 torch.compile 的 dynamic=True ​

8.2 常见错误和解决方案 ​

8.3 调试动态 Shape 问题 ​

升华 ​

"AI 可查 vs 必须理解"清单 ​

动态 Shape——符号分析与形状处理 / Dynamic Shapes, Symbolic Analysis, and Shape Processing

第1节：动态 Shape 的挑战

1.1 为什么动态 Shape 困难

1.2 动态 Shape 的典型场景

1.3 动态 Shape 的类型

第2节：Symbolic Shape——符号化维度

2.1 符号维度的表示

2.2 约束传播

2.3 Symbolic Shape 的实际应用

第3节：Partial Shape——部分已知形状

3.1 Partial Shape 的表示

3.2 Partial Shape 的约束

第4节：Shape Specialization——形状特化

4.1 编译期特化

4.2 JIT 特化

第5节：Dynamic Dimension Bounds——动态维度边界

5.1 为什么要 Bounds

5.2 Bounds 传播

5.3 Bounds 的使用场景

第6节：动态 Shape 下的优化限制

6.1 无法做的优化

6.2 权衡策略

第7节：JAX/XLA 的动态 Shape 支持

7.1 JAX 的 Polymorphic Shapes

7.2 Triton 的动态 Shape 支持

第8节：PyTorch 2.0 的动态 Shape 处理

8.1 torch.compile 的 dynamic=True

8.2 常见错误和解决方案

8.3 调试动态 Shape 问题

升华

"AI 可查 vs 必须理解"清单