Module tree_shape 

Spine+hedge draft-tree shapes for tree speculative decoding (Phase 1 of the tree-spec plan).

A tree is a spine (the drafter’s top-1 chain, exactly today’s chain draft) plus hedge leaves: rank-2..k siblings of spine nodes with no children. The constraint is load-bearing — every node’s ancestors are spine nodes, so per-row attention visibility is “committed prefix + spine(1..d-1) + self”, the drafter KV stays a linear chain, and GDN verification decomposes into one spine pass + one 1-token pass per hedge.

Shape notation (ATLAS_TREE_SHAPE="1,2,2,2"): per-depth node counts; depth d contributes 1 spine node + (count_d - 1) hedges. Verify width M = 1 (root row) + total nodes.

Row layout (fixed): row 0 = root (last committed token, the bonus row), rows 1..=L = spine in depth order, then hedges in (depth, rank) order. The spine being a contiguous prefix is what lets the GDN layer run one existing wy-kernel pass over rows 0..=L unchanged.

Structs

HedgeNode

One hedge leaf: the drafter’s rank-rank candidate at depth.

TreeDraft

A proposed draft tree: spine tokens + hedge leaves for one verify step.

TreeRow

One verify row: (token, depth, parent_row). Row 0 is the root.

TreeShape

Static tree shape: per-depth node counts.

Constants

MAX_NODES

Maximum tree nodes (excl. root row); M = nodes + 1 <= 8 keeps every projection on the batched-GEMV path once the batch8 kernel lands.

MAX_RANK

Maximum candidate rank per position (drafter shadow/top-k width).