Skip to content

Commit bd7534d

Browse files
apstenku123apstenku123
committed
docs: write low-level path_c full-graph runtime blockers contract
The earlier 'Path C: suffix-bypass via mx.custom_function...' chain advertised 'works on local_gb10_quarter' while the kernel was actually consuming all-zero hidden_entry/target_ids/target_mask and emitting bogus grads that silently overrode eager autograd for in-region parameters. Commit b04fbbb removed the unsafe overlay and forced the runtime back to warmup+eager when the suffix loss ABI advertises pending. This commit documents exactly what TileLang/TVM-side work is left so the next pass can actually beat path_b instead of repeating the gate-flipping pattern. This is documentation only; runtime behaviour is unchanged.
1 parent b04fbbb commit bd7534d

1 file changed

Lines changed: 200 additions & 0 deletions

File tree

‎docs/upstream/_path_c_full_graph_runtime_blockers.md‎

Lines changed: 200 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,200 @@
1+
# Path C full-graph runtime — concrete blockers as of 2026-05-23
2+
3+
Status: **does not beat path_b yet.** This document is the lowest-level
4+
contract for the next iteration so any further hand-waving stops here.
5+
6+
## Goal restated
7+
8+
User-stated objective: "полный graph path c через доработку
9+
tilelang -> tvm -> tvm-ffi и чтоб у нас была структура в которую мы
10+
добавляем блоки, с ней tilelang оптимизирует fusion, schedulers,
11+
parallel execution... наш path c должен начать работать быстрее и
12+
экономичнее чем path b. Предпочтение нижнему уровню доработкам на
13+
c/c++/mm/llvm и python в tilelang/tvm/tvm-ffi, и только в крайнем
14+
случае через nanobind и минимальные изменения cppmega.mlx."
15+
16+
Achievement criterion: `scripts/bench_1b_training_matrix.py --paths
17+
path_b,path_c_warm` shows path_c_warm tok/s >= path_b on >= 6/7
18+
optimizers in bf16, with peak GB <= path_b + 10%.
19+
20+
## What is actually wired right now
21+
22+
* `cppmega_mlx/runtime/path_c_fusion.py` builds a `PathCFusionRegion`
23+
from selected route bricks 10 (M) / 11 (R) / 12 (A).
24+
* `cppmega_mlx/runtime/path_c_fusion_schedules.py` lowers that region
25+
to one TileLang/TVM PrimFunc via `plan_path_c_fusion_schedule_for_region`
26+
with `include_backward=True`, plus a suffix loss block.
27+
* `HybridTinyLM.path_c_fused_train_block_prim_func` materialises that
28+
PrimFunc and exposes the physical ABI map (`_cppmega_path_c_physical_buffer_abi_map`)
29+
+ bank shapes. Banks are 3 dtype-uniform arrays (float32 / uint8 / int32),
30+
total ~273 KB for `local_gb10_quarter` tiny smoke at seq_len 127.
31+
* `HybridTinyLM.make_path_c_physical_abi_bank_owner` zero-inits banks
32+
with the right dtypes/shapes.
33+
* `HybridTinyLM.bind_path_c_in_region_parameter_views_into_bank` writes
34+
each in-region parameter into its bank slot (one slice-assign each,
35+
no large staging tensor) and replaces the model attribute with a
36+
bank-view snapshot (`logical_bank_view`).
37+
* `HybridTinyLM.sync_path_c_in_region_parameters_into_bank` re-writes
38+
optimizer-replaced parameter tensors into the bank slots once per
39+
training step.
40+
* `cppmega_mlx/training/path_c_fused_suffix.py` builds an
41+
`mx.custom_function` that on forward writes
42+
(`hidden_entry`, `target_ids`, `target_mask`, *params*) into bank
43+
slots and runs `artifact.forward(bank_owner=...)`, and on backward
44+
returns bank-view cotangents for every primal.
45+
* `PathCFusedPlusEagerTrainingRuntime` honors three modes:
46+
suffix-bypass (only when its install gate accepts), an explicit
47+
fail-closed warmup+eager mode (today's default), and the older
48+
warmup-only path when no aliases are bound.
49+
* `scripts/m04_train_step.py::install_path_c_fused_train_block_runtime_for_model`
50+
now refuses to attach the suffix-bypass loss function when the
51+
generated suffix loss ABI says `"full loss codegen is pending"` or
52+
when scalar outputs are not marked computed.
53+
* `cppmega_mlx/runtime/path_c_physical_abi.py::write_into_bank_slot`
54+
fails closed on dtype mismatch instead of silently coercing.
55+
56+
## What blocks `path_c_warm > path_b`
57+
58+
These are all in the lowest layer of the stack
59+
(`cppmega_mlx/runtime/path_c_fusion_schedules.py` brick descriptors +
60+
TileLang/TVM lowering on the Metal target). cppmega_mlx app code is
61+
intentionally not the place to fix any of them.
62+
63+
### Block 1: brick descriptors do not write all forward outputs
64+
65+
Verified with a direct kernel call against pre-populated bank inputs
66+
(see `tmp/parity_natural.py`):
67+
68+
* `local_gb10_quarter_brick_11_R_hidden_after` slot does receive
69+
data after the kernel runs (sum != 0).
70+
* `local_gb10_quarter_brick_12_A_sparse_mla_fp8_apply_out` stays at
71+
zero even when we supply a non-trivial `hidden_entry` plus
72+
`sparse_mla_sm_scale = 1/sqrt(head_dim)` and `has_sinks = 0`.
73+
74+
That means the row-phased sparse MLA apply descriptor
75+
(`_append_row_phased_sparse_mla_fp8_apply_body` ~line 4888 of
76+
`path_c_fusion_schedules.py`) either skips writing
77+
`attention_out` for this region wiring, or its required FP8/index
78+
inputs (`q_fp8`, `q_scale`, `kv_fp8`, `kv_scale`, sparse `indices`)
79+
are not produced by the upstream qkv_projection descriptor and remain
80+
zero. Until apply writes a non-zero output, suffix loss runs against
81+
half the residual stream and final_norm/lm_head grads diverge from
82+
eager by several percent.
83+
84+
Fix layer: TileLang descriptor for `sparse_mla_fp8_apply` in the
85+
brick-graph schedule, and the qkv_projection descriptor that feeds it.
86+
87+
### Block 2: gradient codegen covers only `final_norm_weight` and
88+
`lm_head_weight`
89+
90+
`_train_step_suffix_loss_parameter_grad_buffers` returns just
91+
`("final_norm_weight_grad", "lm_head_weight_grad")` and
92+
`_append_train_step_suffix_loss_parameter_grads` only writes those
93+
two grad slots. The ABI map advertises 27 grad slots (all in-region
94+
brick params plus norm/head), but 25 of them are left as
95+
uninitialised bank memory by the kernel. Overlaying them onto the
96+
eager grad tree (the original "merged-grad" mode) is therefore unsafe
97+
and currently fail-closed in the runtime.
98+
99+
Fix layer: TileLang descriptors for mamba3 / m2rnn / sparse_mla_fp8_apply
100+
backward, plus matching emit logic in `_append_train_step_suffix_loss_parameter_grads`
101+
so every brick weight grad is generated alongside `final_norm_weight_grad`
102+
and `lm_head_weight_grad`.
103+
104+
### Block 3: suffix loss ABI flag still advertises pending
105+
106+
`_TRAIN_STEP_SUFFIX_LOSS_INPUT_ABI_REASON` literally contains
107+
`"full loss codegen is pending"`. The m04 installer reads this and
108+
declines to attach the suffix-bypass loss function. This is the
109+
correct behaviour: until Block 1 and Block 2 are closed, suffix-bypass
110+
silently trains a different loss. Once Blocks 1 and 2 are closed, the
111+
reason string in
112+
`cppmega_mlx/runtime/path_c_fusion_schedules.py` line ~133 must be
113+
flipped to a non-"pending" value.
114+
115+
### Block 4: `path_c_training_sequence_length(args)` returns
116+
`seq_len - 1`, the bank/PrimFunc are specialised on that, but the
117+
public `model.path_c_fused_in_region_parameter_bank_aliases()` (no
118+
arg) returns offsets for the default 512-length PrimFunc. External
119+
inspectors who call the public helper get offsets that do not match
120+
the installed bank.
121+
122+
Fix layer: thread the install-time sequence_length through
123+
`HybridTinyLM` so the public alias getter returns the bound offsets
124+
(or cache the alias map on the runtime's contract surface and
125+
deprecate the no-arg form).
126+
127+
### Block 5: per-call MLX `value_and_grad` cotangent assumption
128+
129+
`fused_suffix.vjp` discards the upstream loss cotangent and returns
130+
unscaled bank-view cotangents. That is correct only when the trainer
131+
runs `nn.value_and_grad(model, lambda model, batch: model.path_c_fused_suffix_loss(batch))`
132+
with no outer scaling. `CompiledPretrainingStep._accumulate_or_update`
133+
later scales the merged grads by `1 / grad_accum_steps`, so the
134+
single-step assumption holds today. But the VJP should still scale by
135+
the runtime cotangent before returning bank views to be future-proof
136+
under e.g. mixed-precision loss scaling. This is a small fix inside
137+
`cppmega_mlx/training/path_c_fused_suffix.py::fused_suffix_vjp`.
138+
139+
## Lower-level work items in order
140+
141+
These items are ordered so each unblocks the next.
142+
143+
1. Finish the sparse_mla_fp8_apply descriptor so it writes
144+
`attention_out` deterministically for the `brick_12_A` region.
145+
That includes wiring qkv_projection_kv_fp8 / kv_scale /
146+
row-phased indices into apply (today they are zeros). Numerical
147+
gate: `local_gb10_quarter_brick_12_A_sparse_mla_fp8_apply_out` is
148+
bit-equal (or float32 within 1e-4) to an eager reference attention
149+
on the same `hidden_after_11`.
150+
151+
2. Add backward emitters for every in-region brick weight, mirroring
152+
`_append_train_step_suffix_loss_parameter_grads` for mamba3,
153+
m2rnn, sparse-MLA-apply, and residual_norm parameters. Numerical
154+
gate: bank-resident grad slot of each in-region parameter is bit-equal
155+
(or float32 within 1e-4) to the eager grad on the same prefix
156+
hidden_entry.
157+
158+
3. Flip the suffix loss ABI reason away from `"pending"` once
159+
(1) and (2) pass the numerical gate. The m04 installer will then
160+
attach the suffix-bypass loss function automatically.
161+
162+
4. Cotangent scaling in `fused_suffix.vjp`.
163+
164+
5. Sequence-length-aware public alias helper on `HybridTinyLM`.
165+
166+
6. Re-run `scripts/bench_1b_training_matrix.py --paths path_b,path_c_warm`
167+
and confirm path_c_warm tok/s >= path_b on >= 6/7 optimizers in
168+
bf16 with peak GB <= path_b + 10%. Update
169+
`docs/production_kernel_routing.md`.
170+
171+
## Current honest end-state
172+
173+
* Route stays green: `m04_path_c_training_route_available`,
174+
`run_path_c_fused_train_block_route`.
175+
* `parameter_bank_residency_active = True`, but
176+
`bank_grad_overlay_active = False` and
177+
`suffix_bypass_available = False`. Eager remains the source of every
178+
gradient so training is correct; Path C does not yet replace any
179+
eager work.
180+
* Tests: `tests/test_path_c_fused_plus_eager_runtime.py` (20),
181+
`tests/test_path_c_fused_suffix_custom_function.py` (3),
182+
`tests/test_hybrid_lm_path_c_physical_abi_bank_owner.py` (8),
183+
`tests/test_path_c_physical_abi.py` (35 incl. dtype-mismatch
184+
guard), `tests/test_m04_train_step.py -k "path_c or
185+
fused_train_block"` (66), `tests/v4/test_fusion_stage_{a..f}.py`
186+
+ `test_fusion_roadmap_gaps.py` (401). All pass on commit `b04fbbb`.
187+
188+
## Why this matters
189+
190+
The earlier Codex turn happily reported
191+
`returns_full_model_grads = True`, `merged_parameter_count = 27`,
192+
and a `b54c348..222bd9f` chain that "flips the gate to ok". A
193+
critical review showed those grads were the kernel's output against
194+
all-zero hidden_entry / target_ids / target_mask, then overlaid onto
195+
the eager grad tree — a silent training corruption that the unit
196+
tests could not catch (they used a fake artifact). Commit `b04fbbb`
197+
removes the unsafe overlay and replaces it with the fail-closed
198+
warmup+eager path. The work above is what is left to actually beat
199+
path_b, and it lives at the TileLang descriptor + Metal codegen
200+
layer, exactly where the user asked the work to happen.

0 commit comments

Comments
 (0)