cake_gdn: deterministic chunked GDN training kernels (--linear-attention-backend loom) for the head-sharded Qwen GDN layer

[yyihuang]

Summary

A bit-deterministic chunk_gated_delta_rule backend for the Qwen GDN models (Qwen3.5 / 3.6 / 3.8, Qwen3-Next), selected with --linear-attention-backend loom, on the head-sharded linear-attention layer main gained in #3609 / #3610 / #3634.

• --linear-attention-backend loom: deterministic forward + backward for chunk_gated_delta_rule (chunk 64, K = V = 128, in-kernel q/k L2 norm, cu_seqlens packing, fla CP state passing) as generated CUDA kernels in miles_plugins/models/gdn_chunk_train/ — register-MMA kernels with an fp32 running state across chunks (the per-chunk h tensor is bf16 as in FLA), bf16 intermediates where FLA stores bf16 and reductions without device atomics, so repeated passes on identical inputs are bit-identical (asserted by the tests below). Built as torch CUDA extensions (torch.utils.cpp_extension.load, cached under TORCH_EXTENSIONS_DIR); by default all stages compile on the first call (MILES_GDN_CHUNK_TRAIN_PREBUILD=0 for lazy per-stage builds). NVIDIA SM100a / SM103a (Blackwell).
• miles_plugins/models/linear_attn.py: gdn_kernel("loom") returns it, and gdn_recurrence admits context parallelism for loom as it does for fla (the kernels reuse fla's CP state passing). Grouped value heads go to the kernel directly, as main already does for fla. Under loom the short conv runs on fla's Triton kernels (loom_conv_backend(), FLA_CONV_BACKEND overrides) instead of main's mix, whose causal-conv1d CUDA backward accumulates dweight with atomics and does not reproduce conv1d.weight's gradient bit for bit.
• No weight-layout, bridge or converter change: main's group-major GDN q/k/v-projection / short-conv row layout (miles/backends/megatron_utils/megatron_to_hf/linear_attn_layout.py; in_proj_qkv for Qwen3.5, in_proj_qkvz for Qwen3-Next) is used as is.

History. The branch originally carried its own head-sharded GDN module (miles_plugins/models/gdn_attention.py, head-interleaved in_proj_qkv / conv1d rows, bridge and converter hooks). Merging main (6a849af; c5d7db9 then picks up main's B200 CI scheduling change dba9001, and 9437f0b moves the two B200 tests from stage-c-8-gpu-b200 to stage-c-4-gpu-b200 with num_gpus=4) removed it: main's head-sharded layer stores the same row order, so only the backend remains. That merge also deletes miles_plugins/mbridge/gdn_layout.py and its import from the bridges; importing it executed miles_plugins/mbridge/__init__.py, which imports every bridge and the real mbridge package, and that is what failed the two stage-a-cpu shards of the pre-merge CI run (one test in each; mbridge is not installed on the GitHub CPU runners). The evidence below marked pre-merge was measured on the earlier branch states and is kept for the record.

Tests

• tests/fast/models/test_linear_attn_kernel_selection.py: kernel-selection tests on CPU (loom routing, unknown-backend error, the KDA / short-conv backend defaults); fla vs loom parity (atol = rtol = 1e-2; dg's atol scaled by its reference max-abs) and bit determinism of the output and all five input gradients on packed varlen and equal-length batches, with one and two value heads per key head (needs an SM100a / SM103a GPU with fla; skips elsewhere).
• tests/e2e/precision/test_qwen_gdn_tp_cp_parity.py (torchrun, stage-c-4-gpu-b200, num_gpus=4): the head-sharded layer with the loom backend under TP=2 and TP=2 × CP=2 (both HF projection layouts) and CP=2 (qwen3_next layout), 5 cases, against the replicated fla reference of tests/fast-gpu/linear_attn_reference.py — output, input gradient and every parameter gradient (TP-sharded ones gathered across TP, all summed across CP) — plus bit determinism of a second identical pass. --backend fla runs the same five parity cases as a control (the second-pass determinism check is loom-only).
• tests/e2e/megatron/test_qwen3_5_4B_gdn_loom.py (stage-c-4-gpu-b200, num_gpus=4): Qwen3.5-4B full-parameter GRPO, TP=2 + SP, loom backend (MILES_GDN_BACKEND=fla for the A/B reference; record / replay modes for training-side determinism checks).

Evidence after merging main (trees 6a849af, 7020a92, 9437f0b; GPU gates on 4x NVIDIA B200, radixark/miles:dev)

All runs used the radixark/miles:dev container (torch 2.13.0+cu130, fla 0.5.2, as printed by the 6a849af gate step; the later steps ran the same image). The GPU gates ran on single 4x NVIDIA B200 nodes (the generated kernels build for sm_100a); all post-merge evidence is from B200, the SM103a (GB300) build was not re-run after the merge, and the GB300 results in the pre-merge section are from the earlier branch states only. The merged tree was measured at two states: the merge commit 6a849af and the head with the short-conv fix 7020a92. The final head 9437f0b adds main's B200 CI-scheduling merge c5d7db9 (tests/ci, .github workflows, docs/developer, .claude/rules and one-line register_cuda_ci(..., num_gpus=N) edits in existing test files; no change under miles/ or miles_plugins/) and moves the two loom tests to stage-c-4-gpu-b200 with num_gpus=4. No GPU gate was re-run on 9437f0b; on that head only the CI registry check, tests/fast/models on CPU (51 passed, 8 skipped), pre-commit and the CPU suites below ran. At the time of writing, GitHub CI for 9437f0b is held for approval (the repository approves fork runs through a run-ci* label); the last GitHub-executed CI is from the pre-merge head 48d3a00, and the CPU-suite row is a cluster-side proxy for it, not a substitute.

gate	tree	result
tests/fast/models/test_linear_attn_kernel_selection.py (the 2 fla-vs-flashqla cases deselected)	6a849af	10 passed, 2 deselected: 6 kernel-selection / env-default tests plus the 4 fla-vs-loom parity + bit-determinism cases (1 and 2 value heads per key head x packed and equal-length batch)
same, with the new short-conv unit test	7020a92	11 passed, 2 deselected (adds test_loom_short_conv_defaults_to_triton_unless_set)
test_qwen_gdn_tp_cp_parity.py --backend fla (control)	6a849af	PASS, 5 cases
test_qwen_gdn_tp_cp_parity.py (loom; TP2 qwen3_5 / qwen3_next, CP2 qwen3_next, TP2xCP2 both families)	6a849af	test FAILED (exit 1) on its bit-determinism assertion; all relative errors within the 2e-2 tolerance on every rank (rank 0's printed per-case errors all <= 5.1e-3, max dx 5.02e-3); a second pass did not reproduce one parameter gradient bit for bit in the two TP=2 CP=1 cases (qwen3_5 on ranks 1 and 3, qwen3_next on ranks 1 and 2), while outputs and input gradients reproduced on every rank. Cause: main's short_conv_backend() selects fla's "mix" backend (Triton forward, causal-conv1d CUDA backward, which accumulates dweight with atomicAdd in csrc/causal_conv1d_bwd.cu) whenever causal-conv1d imports and FLA_CONV_BACKEND is unset, as in this image; the diagnostic row below pins the gap to that backward
same	7020a92	PASS, 5 cases, every output / input gradient / parameter gradient reproduced bit for bit by the second pass on every rank (rank 0's printed relative errors all <= 5.0e-3, max dx 4.94e-3; every rank within the 2e-2 tolerance)
same with FLA_CONV_BACKEND=mix (diagnostic, pre-fix behaviour)	7020a92	numerics within tolerance, but the second pass differs in conv1d.weight's gradient alone (max abs gap 7.8e-3 on qwen3_next TP2, 1.2e-4 on qwen3_5 TP2): consistent with the mixed short-conv backward being the sole source of the pre-fix nondeterminism (the 6a849af run reported exactly one unreproduced parameter gradient per failing case, the fixed tree reproduces every gradient with fla's Triton conv backward, and forcing the mixed backend brings the gap back in that one parameter)
tests/fast-gpu/test_linear_attn_head_sharded.py (main's own test, 2 ranks)	6a849af	PASS
tests/fast-gpu/test_linear_attn_layer.py (main's own test, 2 ranks)	6a849af	PASS
tests/e2e/megatron/test_qwen3_5_4B_gdn_loom.py, live loom run, 2 rollouts (fresh torch_dist conversion with the merged converter)	6a849af	exit 0; linear_attention_backend = loom in the Megatron args; check_weight_update_equal on (via --ci-test, visual.* skipped): the SGLang-side WeightChecker ran its snapshot / reset / compare cycle on both engines after the weight update without a mismatch; per step (0-3): train_rollout_logprob_abs_diff 0.0083 / 0.0096 / 0.0084 / 0.0072, train_rollout_kl 3.6e-4 / 4.2e-4 / 3.7e-4 / 3.1e-4 (same level as the pre-merge table below); the only error text in the log is SGLang's non-fatal post-warmup freeze_gc failed self-call at server start (also present in the record run); no fallback messages
same test, live loom run (the record run of the determinism round below; reuses the torch_dist checkpoint converted in the 6a849af gate)	7020a92	exit 0; per step (0-3): train_rollout_logprob_abs_diff 0.0084 / 0.0096 / 0.0087 / 0.0073, train_rollout_kl 3.6e-4 / 4.2e-4 / 3.8e-4 / 3.2e-4; grad norm 0 / 0.416 / 0.370 / 0 (non-zero in one of the two optimizer steps of each rollout; rollout/raw_reward 0.016 / 0.109; the logs carry no per-step advantage statistic)
record / replay bit determinism (loom replay A vs B on the recorded batches)	7020a92	one live loom run recorded its 2 rollout batches; two independent --debug-train-only replays of them with loom (separate Ray jobs): per-step grad norm BIT-IDENTICAL over the 4 optimizer steps (max abs diff 0; two of the steps have grad norm 0, the other two 0.416 and 0.370), train data of the saving ranks (TP-rank 0 of each DP group, ranks 0 and 2) BIT-IDENTICAL (torch.equal; 900 tensors equal, 0 differ over 4 rollout x rank files)
CI CPU suites stage-a-cpu x 4 shards + stage-b-cpu, mbridge masked (GitHub CPU-runner proxy) and present	9437f0b vs origin/main dba9001, same node / image / environment	No new failure. Both trees ran the 10 shard-passes in the same container on the same GPU node with the GPUs hidden (the image's test stack imports transformer_engine, which dlopens libcuda at import, so a CPU-only node cannot collect the suites). Per shard-pass FAILED sets: 4 identical (masked a0 8 = 8, masked a2 2 = 2, plain a3 1 = 1, plain b 0 = 0), 5 differ by exactly one test id (masked a1, masked a3, plain a1, plain a2: one extra on this tree; masked b: one extra on origin/main), plain a0 see below; ERROR sets identical except masked a1 (119 extra ERRORs on this tree, all tests/fast/charts tests). The failures shared by both trees are the 7 tests/ci/test/test_build_wheels.py cases, the real-Ray ..._with_their_gpu_slice placement test (an artifact of hiding the GPUs), test_sglang_anthropic_conversion_contract, test_workplace_launch_uses_the_configured_backend, the dual-stack socket test, the secret-flag classification test, the 4 test_native_indexer_export_preserves_projected_rope_and_scores[*-mbridge] cases and test_the_bridge_opts_out_of_the_group_major_names under the mbridge mask, plus the test_kimi_k3_mbridge.py / test_fp32_param_utils.py / test_mbridge_indexer_import collection errors under the mask. The 5 differing test ids (router session tests test_full_model_result_reaches_render_backend_and_snapshot[v1] ConnectionError and test_a_new_root_may_choose_again[v1] HTTP 502; helm chart tests test_grants_nothing_to_a_release_no_platform_client_lives_in and test_the_shared_helpers_render_the_helm_values, which share the chart directory with the 119 helm dependency build ERRORs produced while the masked and present passes ran concurrently in one tree; the origin/main-only test_sampling_params_passed_through[old_sglang_rollout] httpx.ReadError) all pass on this tree when the four affected files run alone, masked and present: 504 passed, 10 skipped, 1 failed in each pass, the one failure being the GPU-slice test above. Plain-pass shard a0 of this tree was re-run alone after its first pass hung in test_subproc_agent_metadata_reaches_reward_model[subproc] while the two passes shared the node: 8 failed, 4972 passed, 19 skipped, exactly origin/main's plain a0 FAILED set (8 failed, 4970 passed, 19 skipped): the 7 test_build_wheels.py cases (jq: command not found in the image) and the GPU-slice test; the hung test passed.
pre-commit run --files on the 11 Python / Markdown files of the change (the generated csrc/ sources and registry_sm_*.json were not passed)	6a849af / 9437f0b	PASS, empty patch / PASS, empty patch (two runs)

Evidence, pre-merge branch states (own head-sharded module)

All runs used the radixark/miles:dev image (torch 2.13.0+cu130, fla 0.5.2, TE 2.17.0), 4 GPUs, on NVIDIA B200 (sm_100a) and NVIDIA GB300 (sm_103a).

Unit / operator / TP x CP (branch state 96b09d6) (pre-merge)

stage	B200	GB300
tests/fast/models/test_gdn_head_sharded_layout.py + tests/fast/test_qwen_gdn_backend.py + tests/fast/backends/megatron_utils/test_fp32_param_utils.py	31 passed (2 pre-existing FlashQLA cases deselected)	31 passed (same)
test_qwen_gdn_tp_cp_parity.py --backend loom (TP2 qwen3_5, TP2 qwen3_next, CP2 qwen3_next, TP2xCP2 qwen3_next; bit-determinism asserted; 4 cases — the merge's rewrite of the test added the fifth case, TP2xCP2 qwen3_5)	4/4 PASS	4/4 PASS
test_qwen_gdn_tp_cp_parity.py --backend fla	4/4 PASS	4/4 PASS
test_qwen3_5_cp_correctness.py CP=2 / CP=4 (existing test on the refactored module)	PASS / PASS	PASS / PASS

Largest relative deviation of the TP-sharded core vs the replicated single-rank fla reference over output, input gradient and every parameter gradient: 1.04e-2 (dhidden, CP=2), everything else <= 9.3e-3, tolerance 2.5e-2. The fla-backend TP2 runs reproduce the reference weight gradients exactly (max abs 0) where no CP reduction is involved.

Qwen3.5-4B end-to-end (tests/e2e/megatron/test_qwen3_5_4B_gdn_loom.py, TP=2 + SP, 2 rollouts, branch state 4434cc8) (pre-merge)

	B200 loom	B200 fla	GB300 loom	GB300 fla
exit	0	0	0	0
--check-weight-update-equal (LM tensors incl. head-interleaved linear_attn.*)	pass	pass	pass	pass
generated GDN extensions built / loaded	12 stages	-	12 stages	-
fallback messages	0	0	0	0
train/train_rollout_logprob_abs_diff (steps 0-3)	0.0083 / 0.0099 / 0.0082 / 0.0071	0.0083 / 0.0096 / 0.0084 / 0.0073	0.0086 / 0.0097 / 0.0085 / 0.0071	0.0082 / 0.0098 / 0.0083 / 0.0073
train/train_rollout_kl (steps 0-3)	3.6e-4 / 4.4e-4 / 3.6e-4 / 3.1e-4	3.6e-4 / 4.2e-4 / 3.7e-4 / 3.2e-4	3.8e-4 / 4.3e-4 / 3.7e-4 / 3.1e-4	3.7e-04 / 4.3e-04 / 3.6e-04 / 3.2e-04

Trainer-vs-inference log-prob agreement is the same with the deterministic backend as with fla. Per-step losses / grad norms of the live runs are not comparable across runs because the sampled rollouts differ (with 8 prompts x 8 samples at 4096 tokens most groups have zero advantage), which is what the replay round below is for.

Training-side bit determinism (record / replay, branch state a6d7601) (pre-merge)

One live loom run per machine (MILES_E2E_MODE=record, 2 rollouts) dumped its rollout batches; the same batches were then trained on with --debug-train-only twice with the loom backend and twice with fla (MILES_E2E_MODE=replay), saving the per-step grad norm and each rank's post-forward train data (log-probs, ref log-probs, advantages, ...). All five stages exit 0 on both machines.

comparison	B200	GB300
loom replay A vs loom replay B, grad norm (4 steps)	BIT-IDENTICAL over 4 steps (max	diff
loom replay A vs loom replay B, train data (900 tensors, 4 rank/rollout files)	BIT-IDENTICAL	BIT-IDENTICAL
fla replay A vs fla replay B, grad norm	BIT-IDENTICAL over 4 steps (max	diff
fla replay A vs fla replay B, train data	BIT-IDENTICAL	BIT-IDENTICAL
loom vs fla on the same batches, grad norm of the one non-zero-advantage step	9.0754 vs 8.1718	7.2589 vs 6.9856
loom vs fla on the same batches, train data	DIFFERENT (644 tensors equal, 256 differ over 4 rollout/rank files)	DIFFERENT (644 tensors equal, 256 differ over 4 rollout/rank files)

The loom backend reproduces every gradient norm and every trainer log-prob bit for bit across independent processes on both architectures. On this TP=2 / SP recipe fla 0.5.2 also replayed bit-identically, so the two backends are compared on equal footing here: their per-token log-probs differ at bf16 level (256 of 900 tensors, the log-prob / ref-log-prob vectors, largest single-token difference 0.88 on B200 and 0.92 on GB300 over 4096-token sequences of a 32-layer model), while both agree with the SGLang inference log-probs to the same degree (train_rollout_logprob_abs_diff ~0.008 in the live runs above).

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[yyihuang]

Performance follow-up: loom vs FLA v0.5.2, and the slow ref_log_probs pass

Environment for everything below: the radixark/miles:dev image (torch 2.13.0+cu130, fla 0.5.2), 4x NVIDIA B200 (sm_100a) and 4x NVIDIA GB300 (sm_103a). Branch state eda2814.

1. Root cause of the slow ref_log_probs pass, and fix (eda2814)

In the first e2e runs the second rollout's perf/ref_log_probs_time was 65 s (B200) / 100 s (GB300) with the loom backend vs ~2-3 s with fla. The extension build directory shows why: the loom stages were compiled lazily on first use, the register-scan variant is selected by (num_seqs, heads), and the first pack with a new sequence count compiled a new stage inside the ref forward (about 45 s of nvcc per stage, serially). Two such builds landed in that rollout's ref window. This is a JIT artefact, not kernel time.

Fix: chunk_gated_delta_rule now triggers a cached, parallel prebuild of all 16 stages for the device architecture on its first call (_jit.prebuild; MILES_GDN_CHUNK_TRAIN_PREBUILD=0 disables it, serial fallback on a concurrent build failure). Cold prebuild of all 16 stages: 86 s on the B200 node, 48 s on the GB300 node; loading from the extension cache afterwards takes 0.3-13 s. With the fix, warm ref_log_probs_time is identical between the backends (table in §3) and no stage compiles after the first call.

2. Operator level: chunk_gated_delta_rule fwd and fwd+bwd, Qwen3.5-4B geometry

Method: CUDA events around one call, 5 warm-up + 20 timed iterations, median; a 512 MiB memset between iterations evicts L2 (outside the timed region). Forward under no_grad; fwd+bwd with requires_grad and backward on a fixed random dO. d_k = d_v = 128, bf16 q/k/v, fp32 gate, use_qk_l2norm_in_kernel=True, packed cu_seqlens. The fla timings include the repeat_interleave of q/k to the value heads that the fla path of the wrapper performs (fla with pre-expanded q/k is 5-15 % faster still and is in the JSON). "Real packs" are greedy packs of the recorded Qwen3.5-4B rollout samples (all 64 samples are 4227-4298 tokens: prompt + 4096-token truncated response) up to 16k / 32k / 64k tokens. Speedup = fla / loom (> 1 means loom is faster). Host wall time of one call equals the GPU time within 0.1 ms on every shape, so there is no host-side launch overhead.

shape (seqs x len, k/v heads)	scan	B200 loom / fla fwd (ms)	B200 fwd x	B200 loom / fla fwd+bwd (ms)	B200 fwd+bwd x	GB300 loom / fla fwd (ms)	GB300 fwd x	GB300 loom / fla fwd+bwd (ms)	GB300 fwd+bwd x
4 x 4096, 8/16 (TP=2 shard)	32r/k2	0.533 / 0.427	0.80	1.785 / 1.922	1.08	0.536 / 0.618	1.15	1.713 / 1.786	1.04
8 x 4096, 8/16	64r/k1	0.917 / 0.635	0.69	3.319 / 2.227	0.67	0.868 / 0.633	0.73	3.082 / 2.101	0.68
16 x 4096, 8/16	64r/k1	1.745 / 1.219	0.70	6.408 / 3.948	0.62	1.606 / 1.164	0.72	5.951 / 3.806	0.64
1 x 16384, 8/16	32r/k2	0.735 / 0.493	0.67	2.505 / 1.852	0.74	0.726 / 0.720	0.99	2.364 / 1.775	0.75
32 x 512, 8/16	64r/k1	0.528 / 0.392	0.74	1.866 / 1.685	0.90	0.526 / 0.566	1.07	1.729 / 1.531	0.89
real pack, 3 x ~4.25k = 12.9k, 8/16	32r/k2	0.477 / 0.435	0.91	1.616 / 1.532	0.95	0.489 / 0.588	1.20	1.504 / 1.836	1.22
real pack, 7 x ~4.3k = 30.1k, 8/16	64r/k1	0.875 / 0.618	0.71	3.174 / 1.979	0.62	0.845 / 0.675	0.80	2.959 / 2.096	0.71
real pack, 15 x ~4.3k = 64.2k, 8/16	64r/k1	1.749 / 1.239	0.71	6.346 / 4.010	0.63	1.609 / 1.182	0.73	5.945 / 3.938	0.66
4 x 4096, 16/32 (TP=1)	64r/k1	0.920 / 0.637	0.69	3.334 / 2.204	0.66	0.875 / 0.614	0.70	3.091 / 2.037	0.66
16 x 4096, 16/32	64r/k1	3.379 / 2.383	0.71	12.327 / 7.672	0.62	3.119 / 2.241	0.72	11.582 / 7.452	0.64
1 x 16384, 16/32	32r/k2	1.056 / 0.719	0.68	3.730 / 2.304	0.62	0.999 / 0.771	0.77	3.477 / 2.512	0.72
real pack 12.9k, 16/32	64r/k1	0.801 / 0.661	0.82	2.883 / 2.161	0.75	0.770 / 0.713	0.93	2.672 / 2.070	0.77
real pack 64.2k, 16/32	64r/k1	3.374 / 2.415	0.72	12.241 / 7.793	0.64	3.106 / 2.280	0.73	11.515 / 7.650	0.66

All 30 shapes, both architectures

case (seqs x len, heads k/v)	T	scan fwd_h/dhu	b200 loom fwd / fla fwd (ms)	b200 fwd x	b200 loom f+b / fla f+b (ms)	b200 f+b x	gb300 loom fwd / fla fwd (ms)	gb300 fwd x	gb300 loom f+b / fla f+b (ms)	gb300 f+b x
4x4096 (16/32)	16384	[64, 1]/[64, 1]	0.920 / 0.637	0.69	3.334 / 2.204	0.66	0.875 / 0.614	0.70	3.091 / 2.037	0.66
8x4096 (16/32)	32768	[64, 1]/[64, 1]	1.755 / 1.219	0.69	6.428 / 3.941	0.61	1.608 / 1.160	0.72	5.959 / 3.820	0.64
16x4096 (16/32)	65536	[64, 1]/[64, 1]	3.379 / 2.383	0.71	12.327 / 7.672	0.62	3.119 / 2.241	0.72	11.582 / 7.452	0.64
1x16384 (16/32)	16384	[32, 2]/[32, 2]	1.056 / 0.719	0.68	3.730 / 2.304	0.62	0.999 / 0.771	0.77	3.477 / 2.512	0.72
32x512 (16/32)	16384	[64, 1]/[64, 1]	0.955 / 0.649	0.68	3.434 / 2.097	0.61	0.906 / 0.647	0.71	3.234 / 2.186	0.68
128x128 (16/32)	16384	[64, 1]/[64, 1]	1.158 / 0.734	0.63	4.161 / 2.616	0.63	1.081 / 0.735	0.68	3.907 / 2.610	0.67
4x4096 (8/16)	16384	[32, 2]/[64, 1]	0.533 / 0.427	0.80	1.785 / 1.922	1.08	0.536 / 0.618	1.15	1.713 / 1.786	1.04
8x4096 (8/16)	32768	[64, 1]/[64, 1]	0.917 / 0.635	0.69	3.319 / 2.227	0.67	0.868 / 0.633	0.73	3.082 / 2.101	0.68
16x4096 (8/16)	65536	[64, 1]/[64, 1]	1.745 / 1.219	0.70	6.408 / 3.948	0.62	1.606 / 1.164	0.72	5.951 / 3.806	0.64
1x16384 (8/16)	16384	[32, 2]/[32, 2]	0.735 / 0.493	0.67	2.505 / 1.852	0.74	0.726 / 0.720	0.99	2.364 / 1.775	0.75
32x512 (8/16)	16384	[64, 1]/[64, 1]	0.528 / 0.392	0.74	1.866 / 1.685	0.90	0.526 / 0.566	1.07	1.729 / 1.531	0.89
128x128 (8/16)	16384	[64, 1]/[64, 1]	0.620 / 0.431	0.69	2.208 / 1.843	0.83	0.608 / 0.599	0.98	2.055 / 1.692	0.82
real 16k_0 (16/32)	12894	[64, 1]/[64, 1]	0.801 / 0.661	0.82	2.883 / 2.161	0.75	0.770 / 0.713	0.93	2.672 / 2.070	0.77
real 16k_0 (8/16)	12894	[32, 2]/[64, 1]	0.477 / 0.435	0.91	1.616 / 1.532	0.95	0.489 / 0.588	1.20	1.504 / 1.836	1.22
real 16k_1 (16/32)	12894	[64, 1]/[64, 1]	0.802 / 0.660	0.82	2.832 / 1.970	0.70	0.771 / 0.741	0.96	2.672 / 2.072	0.78
real 16k_1 (8/16)	12894	[32, 2]/[64, 1]	0.476 / 0.436	0.92	1.566 / 1.502	0.96	0.483 / 0.594	1.23	1.504 / 1.772	1.18
real 16k_2 (16/32)	12855	[64, 1]/[64, 1]	0.792 / 0.657	0.83	2.809 / 1.981	0.71	0.767 / 0.743	0.97	2.656 / 2.070	0.78
real 16k_2 (8/16)	12855	[32, 2]/[64, 1]	0.469 / 0.436	0.93	1.550 / 1.508	0.97	0.486 / 0.620	1.28	1.507 / 1.759	1.17
real 32k_0 (16/32)	30086	[64, 1]/[64, 1]	1.678 / 1.196	0.71	6.081 / 3.886	0.64	1.549 / 1.139	0.74	5.554 / 3.862	0.70
real 32k_0 (8/16)	30086	[64, 1]/[64, 1]	0.875 / 0.618	0.71	3.174 / 1.979	0.62	0.845 / 0.675	0.80	2.959 / 2.096	0.71
real 32k_1 (16/32)	29852	[64, 1]/[64, 1]	1.664 / 1.185	0.71	6.029 / 3.857	0.64	1.531 / 1.129	0.74	5.511 / 3.829	0.69
real 32k_1 (8/16)	29852	[64, 1]/[64, 1]	0.870 / 0.612	0.70	3.105 / 1.956	0.63	0.827 / 0.650	0.79	2.923 / 2.078	0.71
real 32k_2 (16/32)	29753	[64, 1]/[64, 1]	1.656 / 1.179	0.71	6.003 / 3.835	0.64	1.524 / 1.128	0.74	5.478 / 3.812	0.70
real 32k_2 (8/16)	29753	[64, 1]/[64, 1]	0.866 / 0.612	0.71	3.096 / 1.952	0.63	0.823 / 0.647	0.79	2.914 / 2.076	0.71
real 64k_0 (16/32)	64197	[64, 1]/[64, 1]	3.374 / 2.415	0.72	12.241 / 7.793	0.64	3.106 / 2.280	0.73	11.515 / 7.650	0.66
real 64k_0 (8/16)	64197	[64, 1]/[64, 1]	1.749 / 1.239	0.71	6.346 / 4.010	0.63	1.609 / 1.182	0.73	5.945 / 3.938	0.66
real 64k_1 (16/32)	63747	[64, 1]/[64, 1]	3.336 / 2.401	0.72	12.141 / 7.723	0.64	3.014 / 2.262	0.75	11.308 / 7.597	0.67
real 64k_1 (8/16)	63747	[64, 1]/[64, 1]	1.735 / 1.230	0.71	6.298 / 3.961	0.63	1.593 / 1.174	0.74	5.898 / 3.921	0.66
real 64k_2 (16/32)	63767	[64, 1]/[64, 1]	3.333 / 2.403	0.72	12.139 / 7.734	0.64	3.013 / 2.263	0.75	11.308 / 7.598	0.67
real 64k_2 (8/16)	63767	[64, 1]/[64, 1]	1.735 / 1.232	0.71	6.299 / 3.960	0.63	1.593 / 1.175	0.74	5.901 / 3.930	0.67

Takeaways: the deterministic kernels are at parity with fla only on the small TP=2 shard packs that the e2e actually runs (<= 16k tokens per micro-batch with --max-tokens-per-gpu 16384, 3-4 sequences, r32k2 scan): 0.9-1.3x fwd+bwd. On >= 32k-token packs and on all 16/32-head (TP=1) shapes loom is slower, fwd 0.68-0.77x and fwd+bwd 0.61-0.72x of fla, on both architectures.

Is that the scan-configuration heuristic? No. Forcing each register-scan configuration on the slow shapes (GB300):

     T seqs  heads       auto     r32k2 f/fb     r32k1 f/fb     r64k1 f/fb       fla f/fb
 32768    8  8/16     (64, 1)   0.942/ 3.385   0.940/ 3.467   0.873/ 3.101  0.548/ 1.718
 65536   16  8/16     (64, 1)   1.748/ 6.449   1.745/ 6.562   1.603/ 5.947  0.817/ 3.000
 30100    7  8/16     (64, 1)   0.912/ 3.136   0.908/ 3.195   0.837/ 2.941  0.568/ 1.775
 16384    4  8/16     (32, 2)   0.537/ 1.789   0.534/ 1.824   0.515/ 1.673  0.504/ 1.565
 16384    1  8/16     (32, 2)   0.726/ 2.370   0.788/ 2.578   0.806/ 2.636  0.655/ 1.590
 16384   32  8/16     (64, 1)   0.567/ 1.884   0.563/ 1.912   0.532/ 1.736  0.490/ 1.312
 32768    8 16/32     (64, 1)   1.759/ 6.481   1.755/ 6.577   1.610/ 5.962  0.796/ 3.028
 65536   16 16/32     (64, 1)   3.428/12.812   3.414/12.942   3.114/11.587  1.446/ 5.754
 30100    7 16/32     (64, 1)   1.571/ 5.844   1.562/ 5.927   1.542/ 5.529  0.834/ 3.134
 16384    4 16/32     (64, 1)   0.942/ 3.396   0.946/ 3.456   0.882/ 3.099  0.580/ 1.738
 16384    1 16/32     (32, 2)   1.004/ 3.482   1.076/ 3.701   1.096/ 3.761  0.695/ 2.182
 16384   32 16/32     (64, 1)   0.994/ 3.567   0.995/ 3.603   0.907/ 3.235  0.580/ 1.817

choose_register_scan_config already picks the fastest variant on every shape (r64k1 where it is chosen, r32k2 for the single long sequence); the remaining gap to fla is in the kernels themselves (the state-scan and the dhu/dqkwg backward stages scale ~1.6-1.7x worse with tokens than fla's Triton kernels). That is a kernel-side item and is tracked in the kernel generator's issue tracker, not in this PR.

3. End-to-end, warm: Qwen3.5-4B GRPO, TP=2 + SP, 4 GPUs colocated

One live loom run recorded 5 rollouts (8 prompts x 8 samples, 4096-token responses); the same batches were then replayed train-only (--debug-train-only) with the loom backend and with fla, all extensions prebuilt, Triton cache warm from the earlier runs. Medians over rollouts 2-4 (rollouts 0-1 carry module loading / Triton autotuning for both backends):

metric (median of rollouts 2-4, replay on identical batches)	b200 loom	b200 fla	b200 loom/fla	gb300 loom	gb300 fla	gb300 loom/fla
actor train (s)	5.91	5.58	0.94x	5.98	6.75	1.13x
ref log-probs (s)	1.50	1.44	0.96x	1.90	1.91	1.00x
log-probs (s)	1.15	1.16	1.01x	1.31	1.39	1.06x
step (s)	22.03	21.49	0.98x	12.95	13.64	1.05x
actor train TFLOPS	263.94	279.72	0.94x	260.31	231.34	1.13x
actor train tok/s	46211.95	48974.66	0.94x	45596.57	40502.63	1.13x

Per-rollout perf/ref_log_probs_time (s), showing the JIT effect and its removal:

• b200 loom live (record, prebuilt): r0=7.96, r1=8.51, r2=1.77, r3=1.80, r4=1.69; actor_train: r0=15.74, r1=16.79, r2=6.07, r3=5.97, r4=6.01
• b200 loom replay: r0=6.08, r1=2.20, r2=1.45, r3=1.84, r4=1.50; actor_train: r0=15.63, r1=6.77, r2=5.76, r3=6.10, r4=5.91
• b200 fla replay: r0=16.10, r1=2.00, r2=1.80, r3=1.44, r4=1.44; actor_train: r0=67.46, r1=46.27, r2=6.06, r3=5.45, r4=5.58
• gb300 loom live (record, prebuilt): r0=22.56, r1=2.32, r2=2.41, r3=1.87, r4=1.78; actor_train: r0=22.84, r1=6.23, r2=6.26, r3=6.03, r4=6.03
• gb300 loom replay: r0=5.70, r1=2.12, r2=1.72, r3=2.05, r4=1.90; actor_train: r0=13.48, r1=6.15, r2=5.98, r3=5.96, r4=6.14
• gb300 fla replay: r0=17.95, r1=2.23, r2=1.91, r3=2.01, r4=1.90; actor_train: r0=83.16, r1=58.91, r2=6.75, r3=6.83, r4=6.65

At the model level the two backends are within +-13 % of each other (loom 0.94x on B200, 1.13x on GB300 on actor_train_time), and the ref/log-prob passes are identical, consistent with the GDN layers being a small fraction of the 4B model's step and with the operator being at parity on the <= 16k-token micro-batches this recipe produces. fla also pays JIT inside training: its rollouts 0 and 1 in the replay show 67 s / 46 s (B200) and 83 s / 59 s (GB300) actor_train_time from Triton compilation, vs 6 s warm.

4. Gating of the code change

On B200 with eda2814: tests/fast/test_qwen_gdn_backend.py 4 passed (2 pre-existing FlashQLA cases deselected); tests/e2e/precision/test_qwen_gdn_tp_cp_parity.py --backend loom TP2 (both HF layouts) / CP2 / TP2xCP2: 4/4 PASS with the loom bit-determinism assertions; the 5-rollout live loom run and both replays exit 0 with zero backend fallbacks.

[yyihuang]

Qwen3.8-27B dense GRPO on this branch (9437f0b)

Recipe: scripts/run_qwen3_dense.py --model-name Qwen3.8-27B unchanged (1 x 8 GPU, TP4, SP, full recompute, CPU Adam, max_tokens_per_gpu 8192, rollout-batch-size 32, n-samples-per-prompt 8, rollout-max-response-len 8192, global-batch-size 256), dapo-math-17k prompts, AIME-2024 eval; only --extra-args "--linear-attention-backend <loom|fla>" differs between the arms. Environment: radixark/miles:dev image (torch 2.13.0+cu130, fla 0.5.2), one 8 x B200 node (192 cores, 2 TiB host RAM), container launched with --no-container-mount-home --container-writable.

Preparation: HF Qwen/Qwen3.8-27B (18 safetensors shards, 52 GB; the shared model catalogs are not mounted on this cluster) downloaded once in 69 s; converted with this branch's tools/convert_hf_to_torch_dist.py + model_args_utils.py qwen3.8-27B in 114 s (latest_checkpointed_iteration.txt = release, 51 GB); the 16 loom extension stages were built once (307 s) and reused by every arm.

Summary

• The documented 27B recipe runs end to end on this branch with --linear-attention-backend loom: 10 rollout steps + AIME eval before and after, no OOM / NaN / training-side traceback, loom selected on all 8 ranks (no fallback). Step 0 rollout/raw_reward 0.6406 and train/train_rollout_logprob_abs_diff 0.0104 reproduce the documented 0.64 / 0.010; the logprob diff stays in 0.0099–0.0109 over all 10 steps (fla control: 0.0105–0.0106), train_rollout_kl 3.6e-4 to 4.0e-4 in both arms, grad norm 0.06–0.16, rewards 0.48–0.73 with no downward trend; AIME-2024 went from 0.5229 (initial weights) to 0.6104 after the 10 steps (480 samples each, so still noisy).
• Steady-state wall time on 8 x B200 is about 5.1 min per step (rollout ≈ 205 s, train ≈ 89 s, weight sync ≈ 3.7 s, 791 rollout tokens/GPU/s) against the documented ≈ 12 min / 478 tokens/GPU/s on 8 x H200.
• loom vs fla on the same node: step-time ratio 0.96 (307 s vs 319 s), train-phase ratio 0.945 (89 s vs 95 s); the fla arm's sampled responses were about 10 % longer, and the length-normalized actor-train throughput is within 5 %, so at 27B with this recipe the two backends are indistinguishable in speed (the dense layers dominate). No speedup is claimed here; the point of the loom arm is the bit-deterministic training below.

loom arm (10 rollout steps)

step	raw_reward	response_lengths	truncated	train_rollout_logprob_abs_diff	train_rollout_kl	loss	grad_norm	rollout_time	ref_log_probs_time	log_probs_time	actor_train_time	train_time	update_weights_time	step_time	tokens_per_gpu_per_sec
0	0.6406	4532	0.3203	0.01043	0.0003789	2.794e-09	0.07667	192.1	17.92	11.49	137.4	184.3	3.396	1858	755
1	0.4844	5382	0.5117	0.0106	0.0003895	-6.985e-10	0.07855	200	15.45	13.57	60.58	90.51	3.78	303.5	861.1
2	0.5273	5472	0.4375	0.01049	0.0003821	-7.451e-09	0.1247	268.1	21.88	13.72	72.52	108.7	3.648	405.7	653.1
3	0.6328	4824	0.3516	0.009855	0.0003631	-5.588e-09	0.08117	179.8	13.56	12.13	56	82.26	4.007	273.1	858.4
4	0.6016	5184	0.3984	0.01089	0.0003986	6.519e-09	0.08336	198.3	14.8	13.09	59.58	88.64	3.637	297.3	836.7
5	0.5	5290	0.457	0.0103	0.0003756	7.567e-09	0.1007	222	15.12	13.31	61.72	90.72	3.887	323.6	762.6
6	0.668	4484	0.332	0.01014	0.0003712	7.683e-09	0.0742	157.4	12.59	11.33	52.68	77.16	3.441	245.3	911.4
7	0.5898	5181	0.4102	0.01059	0.0003869	-4.657e-09	0.0747	218.2	15.12	13.28	70.49	105.8	3.764	334.8	760
8	0.7305	4008	0.2656	0.009984	0.0003735	1.863e-08	0.09433	205.2	12.01	10.34	48.98	71.89	3.923	287.9	625.1
9	0.5938	5146	0.3828	0.01008	0.0003809	5.588e-09	0.1606	192.9	14.32	13.07	59.87	87.84	3.641	291.5	853.4

Evaluation (AIME-2024, 30 problems x 16 samples, loom arm; weight_version v means the weights after v - 1 optimizer steps, so eval 0 used the initial weights and eval 9 the weights after all 10 steps):

• eval 0: aime = 0.5229, aime/num_training_samples = 480, aime/response_len/mean = 1.05e+04, aime/truncated_ratio = 0.475, aime/weight_version/mean = 1
• eval 9: aime = 0.6104, aime/num_training_samples = 480, aime/response_len/mean = 9948, aime/truncated_ratio = 0.3833, aime/weight_version/mean = 11

fla control arm (3 rollout steps)

step	raw_reward	response_lengths	truncated	train_rollout_logprob_abs_diff	train_rollout_kl	loss	grad_norm	rollout_time	ref_log_probs_time	log_probs_time	actor_train_time	train_time	update_weights_time	step_time	tokens_per_gpu_per_sec
0	0.6445	4577	0.332	0.01053	0.0003828	-1.281e-09	0.06404	195.7	45.46	11.63	161.5	219.2	3.495	449.8	748.3
1	0.4883	5409	0.5039	0.01058	0.0003891	2.654e-08	0.09716	201.2	15.64	13.69	61.24	91.17	3.916	303	860.5
2	0.5273	5575	0.4414	0.01051	0.0003865	-5.588e-09	0.09161	227	17.47	15.5	64.36	97.86	3.351	335	786.1

loom vs fla (same recipe, same machine)

metric	loom	fla	loom / fla
step time (mean over steps >= 1)	307	319	0.962
rollout time (mean over steps >= 1)	204.7	214.1	0.956
train time (ref + actor log-probs + actor train) (mean over steps >= 1)	89.28	94.51	0.945
ref log-probs time (mean over steps >= 1)	14.98	16.55	0.905
actor log-probs time (mean over steps >= 1)	12.65	14.6	0.867
actor train time (mean over steps >= 1)	60.27	62.8	0.960
weight sync time (mean over steps >= 1)	3.748	3.634	1.031
rollout tokens/GPU/s (mean over steps >= 1)	791.3	823.3	0.961
actor train tokens/s (length-normalized) (mean over steps >= 1)	2.216e+04	2.318e+04	0.956
mean response length (sampled) (mean over steps >= 1)	4997	5492	0.910
logprob abs diff (mean over steps >= 1)	0.01032	0.01054	0.979
train-rollout KL (mean over steps >= 1)	0.0003802	0.0003878	0.980

Bit determinism (record one rollout with loom, replay it train-only twice)

### DET (02:45:10): record 1 rollout with loom, replay twice train-only, compare
ARM_DET_RECORD_EXIT=0 (14min) log=arm_det_record.log
rollout dumps: 1
ARM_DET_REPLAY_A_EXIT=0 (8min) log=arm_det_replay_a.log
ARM_DET_REPLAY_B_EXIT=0 (8min) log=arm_det_replay_b.log
### DET_COMPARE (03:15:39)
a vs b: MISSING
train data a vs b: BIT-IDENTICAL (1794 tensors equal, 0 differ over 2 rollout/rank files)
DET_COMPARE_EXIT=0 (2s)
### DET_COMPARE_LOGS (03:15:41)
train/loss: IDENTICAL (2 vs 2 values) ['6.752088665962219e-09', '6.752088665962219e-09'] vs ['6.752088665962219e-09', '6.752088665962219e-09']
train/pg_loss: IDENTICAL (2 vs 2 values) ['6.752088665962219e-09', '6.752088665962219e-09'] vs ['6.752088665962219e-09', '6.752088665962219e-09']
train/entropy_loss: IDENTICAL (2 vs 2 values) ['0.0', '0.0'] vs ['0.0', '0.0']
train/grad_norm: IDENTICAL (2 vs 2 values) ['0.07266654819250107', '0.07266654819250107'] vs ['0.07266654819250107', '0.07266654819250107']
train/kl_loss: IDENTICAL (2 vs 2 values) ['0.0', '0.0'] vs ['0.0', '0.0']
DET_COMPARE_LOGS_EXIT=0 (0s)
=== DONE Fri Oct  9 03:15:41 AM PDT 2026 rc=0 total=153min ===

a vs b: MISSING refers to the --ci-save-grad-norm files: in this tree the trainer only writes them under --ci-test (check_grad_norm is called from the Megatron train step only when args.ci_test is set), which the documented launcher does not pass, so no grad-norm .pt files exist and the grad norm is compared through the logged values instead (identical to the last printed digit in the recording run and in both replays: grad norm 0.07266654819250107, loss 6.752088665962219e-09, train_rollout_logprob_abs_diff 0.010454539209604263, train_rollout_kl 0.00037983717629686). The train tensors saved by --save-debug-train-data (1794 tensors over the two data-parallel shard files, ranks 0 and 4 of the TP4 layout) are bit-identical between the two replays.

Caveats

• step_time of step 0 includes the pre-training AIME eval (loom arm) and the wait for the first weight load; compare steady-state steps (>= 1) only.
• The H200 numbers in docs/models/qwen/qwen3-8.md (rollout 307 s, train 415 s, weight sync 2.5 s, about 12 min per step, 478 tokens/GPU/s) were measured on different hardware with the default backend and are quoted for orientation only; the controlled comparison is the loom-vs-fla table above (same node, same prepared checkpoint, same recipe, back to back).
• Rollouts are sampled, so the two arms see different trajectories; rewards and response lengths are expected to differ step by step, and 10 (loom) / 3 (fla) steps are too few to read a reward trend. The logprob-diff and timing columns are the comparable ones.
• Host memory: CPU Adam for the 27B model holds about 113 GiB of optimizer state per training rank; a 512 GiB memory cap OOM-killed the first optimizer step of a 4-GPU smoke run, so the 8-GPU runs used a 1500 GiB cap (peak observed 926 GiB of 2015).
• Loom backend evidence: the parsed-args dump shows linear_attention_backend loom, every rank logs Linear-attention short conv backend under loom: triton, and no fallback message appears; the only Traceback blocks in the launcher logs are sglang's _freeze_gc_after_server_warmup readiness polls (connection refused before the HTTP ports open) and occur identically in the fla arm.
• The torch_dist conversion used 4 ranks (--nproc-per-node 4) on a 4-GPU preparation allocation instead of the documented 8; the saved checkpoint is rank-count independent and the TP4 x 8-GPU runs loaded it without changes.

🤖 Generated with Claude Code

[yyihuang]

Qwen3.8-27B dense GRPO on this branch (9437f0b + updated gdn_chunk_train plugin kernels)

Recipe: scripts/run_qwen3_dense.py --model-name Qwen3.8-27B unchanged (1 x 8 GPU, TP4, SP, full recompute, CPU Adam, max_tokens_per_gpu 8192, rollout-batch-size 32, n-samples-per-prompt 8, rollout-max-response-len 8192, global-batch-size 256), dapo-math-17k prompts, AIME-2024 eval; only --extra-args "--linear-attention-backend <loom|fla>" differs between the arms. Environment: radixark/miles:dev image (torch 2.13.0+cu130, fla 0.5.2), one 8 x B200 node (192 cores, 2 TiB host RAM), container launched with --no-container-mount-home --container-writable; same image and node type as the previous comment.

Preparation: The prepared HF Qwen/Qwen3.8-27B checkpoint, its torch_dist conversion and the datasets (dapo-math-17k prompts, AIME-2024 eval) from the previous comment were reused unchanged. Before any arm ran, the plugin directory miles_plugins/models/gdn_chunk_train of this PR's tree (head 9437f0b) received the kernel sources and registry of the current loom GDN export (loom commit ed1d12ef3ad, sm_100a: 44 generated source files; ops.py, _jit.py and the wiring test are the PR-head files, unchanged). These kernels are the current state of the deterministic chunk kernels and differ from the plugin kernels in the PR head; they are staged for a follow-up commit and are not part of this PR yet. The 22 extension modules were prebuilt once (146 s) and the wiring test passed (43 tests). The fla control arm uses the same tree; only --linear-attention-backend differs.

Summary

• The documented 27B recipe runs end to end with --linear-attention-backend loom and the updated plugin kernels: 10 rollout steps + AIME eval before and after, no OOM / NaN / training-side traceback, loom selected on all 8 ranks (no fallback). Step 0 rollout/raw_reward 0.6328 and train/train_rollout_logprob_abs_diff 0.0105 reproduce the documented 0.64 / 0.010; the logprob diff stays in 0.0096–0.0109 over all 10 steps (fla control: 0.0104–0.0105), train_rollout_kl 3.5e-4 to 4.0e-4 in both arms, grad norm 0.07–0.16, rewards 0.50–0.75 with no downward trend; AIME-2024 went from 0.5271 (initial weights) to 0.6417 after the 10 steps (previous comment: 0.5229 to 0.6104; 480 samples per eval, so still noisy).
• Steady-state wall time on 8 x B200 is about 4.9 min per step (rollout ≈ 200 s, train ≈ 81 s = ref log-probs 13.4 + actor log-probs 11.9 + actor train 54.5 + wait, weight sync ≈ 3.5 s, 792 rollout tokens/GPU/s); the previous comment measured ≈ 5.1 min (rollout ≈ 205 s, train ≈ 89 s) with the PR-head kernels on the same node type.
• loom vs fla, sampled arms on the same node: step-time ratio 0.984 (295.5 s vs 300.2 s), train-phase ratio 0.903 (81.0 s vs 89.7 s) — but the fla arm sampled about 11 % longer responses, and the length-normalized actor-train throughput is identical (23,887 vs 23,891 tokens/s; previous comment 22,160 vs 23,180). The controlled comparison is the same-data replay below: the same three recorded rollouts trained train-only with each backend. There the actor-train phase is faster with loom on every step (74.99 / 58.45 / 60.99 s and 75.43 / 58.42 / 60.69 s in the two loom replays vs 79.70 / 59.95 / 61.74 s with fla; mean 64.8 vs 67.1 s, ratio 0.966), the actor log-probs pass is faster on every step (mean 12.8 vs 13.4 s) and the ref log-probs pass is faster on steps 1–2 (14.7 / 15.6 s and 14.5 / 15.1 s vs 14.9 / 15.7 s; step 0 is the first pass and differs more between the two identical loom replays, 38.4 vs 18.4 s, than between backends). Total train_time on steps 1–2: loom 87.3 / 91.0 s and 86.8 / 89.95 s vs fla 89.3 / 125.4 s; the fla step-2 value contains about 30 s outside the three timed phases and the recorded wait (its actor train 61.7 s and log-prob passes are normal), so the like-for-like train-phase gain is 2–4 %, consistent with the GDN kernels being a small share of a dense-27B step. Kernel share, same method as the previous comment (the recorded rollout's production micro-batch schedule, 76 micro-batches per data-parallel rank x 48 GDN layers, replayed through both backends' kernels on one B200: median of 20 L2-flushed calls, a back-to-back loop, and CUPTI device time): per training step (data-parallel rank 0; rank 1 agrees within 0.2 s) the GDN kernels take 4.6 s with loom vs 9.8 s with fla by single-call wall-clock (5.7 % vs 11.0 % of the train phase; two forward-only passes 1.56 vs 3.63 s, train pass 3.00 vs 6.16 s), 4.1 vs 8.2 s back-to-back, and 2.86 vs 4.29 s of device time (3.6 % vs 4.8 %); per micro-batch fla/loom median 2.02 (wall f+b) and 1.42 (device f+b), loom faster by wall-clock in 152 of 152 micro-batches, slower by device time in 4 of 152 (f+b). With the PR-head kernels the previous comment measured loom at 8.6–9.1 s wall / 7.0 s device per step (9.8 % / 8.0 %), i.e. the updated kernels roughly halve loom's kernel wall time and cut its device time by about 59 %, moving it below fla's (previous fla reference 11.2–14.2 s wall / 5.1 s device on that rollout). Shares use the sampled arms' steady-state phase means as denominators and are indicative (about ±1 pp). bf16 forward parity loom vs fla on the real micro-batches: max |diff| 4.9e-4.
• Bit determinism: the two loom replays of the same three recorded rollouts produced bit-identical training tensors (5,382 tensors over 6 rollout/rank dump files, 0 differ) and identical loss / pg_loss / entropy / grad-norm / kl values in the logs; details below.

loom arm (10 rollout steps)

step	raw_reward	response_lengths	truncated	train_rollout_logprob_abs_diff	train_rollout_kl	loss	grad_norm	rollout_time	ref_log_probs_time	log_probs_time	actor_train_time	train_time	update_weights_time	step_time	tokens_per_gpu_per_sec
0	0.6328	4646	0.3359	0.01047	0.0003808	0	0.09202	196.9	18.19	11.37	90.39	126.3	3.406	1626	755.2
1	0.5	5323	0.4883	0.0106	0.0003889	-3.725e-09	0.09622	193	14.41	12.89	57.72	85.75	3.46	290.7	882.7
2	0.5273	5580	0.4336	0.01061	0.0003891	1.63e-08	0.1042	258.6	15	13.48	61.32	95.85	3.063	391.5	690.6
3	0.6133	4823	0.375	0.009767	0.000359	1.071e-08	0.07589	173.5	13.07	11.79	53.69	79.12	3.796	265	889.8
4	0.5664	5182	0.4258	0.01086	0.0004	2.328e-09	0.09317	184.5	13.99	12.6	56.09	83.32	3.504	279	899.1
5	0.5156	5129	0.4336	0.01022	0.0003736	-1.583e-08	0.1075	238.1	13.84	12.48	57.32	84.68	3.728	334.3	689.5
6	0.7031	4305	0.3008	0.009893	0.0003691	0	0.07325	168.1	11.87	10.46	49.23	72.14	3.445	251.5	819.3
7	0.5703	5166	0.4141	0.01037	0.00038	1.723e-08	0.1092	229.7	14.14	12.49	56.42	83.6	3.412	324.9	719.7
8	0.7539	3712	0.2266	0.009551	0.0003526	0	0.1155	162.5	10.63	9.07	44.21	64.44	3.762	238.5	730.9
9	0.6484	4843	0.3242	0.009647	0.0003553	9.313e-10	0.163	193.2	13.41	11.63	54.19	79.99	3.672	284.2	802.2

Evaluation (AIME-2024, 30 problems x 16 samples, loom arm; weight_version v means the weights after v - 1 optimizer steps, so eval 0 used the initial weights and eval 9 the weights after all 10 steps):

• eval 0: aime = 0.5271, aime/num_training_samples = 480, aime/response_len/mean = 1.055e+04, aime/truncated_ratio = 0.4771, aime/weight_version/mean = 1
• eval 9: aime = 0.6417, aime/num_training_samples = 480, aime/response_len/mean = 9478, aime/truncated_ratio = 0.3438, aime/weight_version/mean = 11

fla control arm (3 rollout steps)

step	raw_reward	response_lengths	truncated	train_rollout_logprob_abs_diff	train_rollout_kl	loss	grad_norm	rollout_time	ref_log_probs_time	log_probs_time	actor_train_time	train_time	update_weights_time	step_time	tokens_per_gpu_per_sec
0	0.6562	4572	0.3164	0.0104	0.0003762	-1.863e-09	0.07095	194	34.55	11.47	157.9	213.9	4.248	464.3	754.2
1	0.4805	5384	0.5195	0.01055	0.0003911	1.956e-08	0.07077	183.2	15.39	13.52	60.06	89.5	3.584	283	940.4
2	0.5547	5472	0.4102	0.0105	0.0003836	1.583e-08	0.09485	217.5	15.07	13.76	60.42	89.84	3.581	317.3	805.3

loom vs fla (same recipe, same machine)

metric	loom	fla	loom / fla
step time (mean over steps >= 1)	295.5	300.2	0.984
rollout time (mean over steps >= 1)	200.1	200.3	0.999
train time (ref + actor log-probs + actor train) (mean over steps >= 1)	80.99	89.67	0.903
ref log-probs time (mean over steps >= 1)	13.37	15.23	0.878
actor log-probs time (mean over steps >= 1)	11.88	13.64	0.871
actor train time (mean over steps >= 1)	54.46	60.24	0.904
weight sync time (mean over steps >= 1)	3.538	3.583	0.988
rollout tokens/GPU/s (mean over steps >= 1)	791.5	872.9	0.907
actor train tokens/s (length-normalized) (mean over steps >= 1)	2.389e+04	2.389e+04	1.000
mean response length (sampled) (mean over steps >= 1)	4896	5428	0.902
logprob abs diff (mean over steps >= 1)	0.01017	0.01052	0.966
train-rollout KL (mean over steps >= 1)	0.0003742	0.0003874	0.966

Bit determinism and same-data timing (record three rollouts with loom, replay them train-only twice with loom and once with fla)

### DET (08:14:44): record 3 rollouts with loom, replay them train-only with loom (a, b) and fla (f), compare;
### ARM det_record (08:14:44): num_rollout=3 --no-enable-eval extra='--linear-attention-backend loom --save-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt'
ARM_DET_RECORD_EXIT=0 (24min)
rollout dumps: 3 (expected 3)
### ARM det_replay_a (08:38:55): num_rollout=3 --no-enable-eval extra='--linear-attention-backend loom --load-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt --debug-train-only --ci-save-grad-norm <workspace>/det/grad_norm_a_{rollout_id}_{step_id}.pt --save-debug-train-data <workspace>/det/train_a_{rollout_id}_{rank}.pt'
ARM_DET_REPLAY_A_EXIT=0 (12min)
### ARM det_replay_b (08:51:05): num_rollout=3 --no-enable-eval extra='--linear-attention-backend loom --load-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt --debug-train-only --ci-save-grad-norm <workspace>/det/grad_norm_b_{rollout_id}_{step_id}.pt --save-debug-train-data <workspace>/det/train_b_{rollout_id}_{rank}.pt'
ARM_DET_REPLAY_B_EXIT=0 (12min)
### ARM det_replay_f (09:03:14): num_rollout=3 --no-enable-eval extra='--linear-attention-backend fla --load-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt --debug-train-only'
ARM_DET_REPLAY_F_EXIT=0 (12min)
### DET_COMPARE (09:15:24)
a vs b: MISSING
train data a vs b: BIT-IDENTICAL (5382 tensors equal, 0 differ over 6 rollout/rank files)
DET_COMPARE_EXIT=0 (2s)
### DET_COMPARE_LOGS (09:15:26)
train/loss: IDENTICAL (6 vs 6 values) ['-2.3283064365386963e-09', '-2.3283064365386963e-09', '1.862645149230957e-09', '1.862645149230957e-09'] vs ['-2.3283064365386963e-09', '-2.3283064365386963e-09', '1.862645149230957e-09', '1.862645149230957e-09']
train/pg_loss: IDENTICAL (6 vs 6 values) ['-2.3283064365386963e-09', '-2.3283064365386963e-09', '1.862645149230957e-09', '1.862645149230957e-09'] vs ['-2.3283064365386963e-09', '-2.3283064365386963e-09', '1.862645149230957e-09', '1.862645149230957e-09']
train/entropy_loss: IDENTICAL (6 vs 6 values) ['0.0', '0.0', '0.0', '0.0'] vs ['0.0', '0.0', '0.0', '0.0']
train/grad_norm: IDENTICAL (6 vs 6 values) ['0.0626189187169075', '0.0626189187169075', '0.061997514218091965', '0.061997514218091965'] vs ['0.0626189187169075', '0.0626189187169075', '0.061997514218091965', '0.061997514218091965']
train/kl_loss: IDENTICAL (6 vs 6 values) ['0.0', '0.0', '0.0003905287594534457', '0.0003905287594534457'] vs ['0.0', '0.0', '0.0003905287594534457', '0.0003905287594534457']
DET_COMPARE_LOGS_EXIT=0 (0s)
=== DONE Sun Oct 11 09:15:26 AM PDT 2026 rc=0 total=178min ===

Per-step timings of the three train-only replays of the same recorded rollouts (seconds; train_time = ref log-probs + actor log-probs + actor train + waits):

replay	step	train_time	ref log-probs	actor log-probs	actor train	train wait	grad_norm	logprob abs diff
a (loom)	0	125.3	38.35	11.33	74.99	6.32	0.06262	0.01047
a (loom)	1	87.28	14.67	13.25	58.45	6.16	0.062	0.01052
a (loom)	2	91.04	15.59	13.93	60.99	6.227	0.09744	0.01047
b (loom)	0	106	18.36	11.31	75.43	1.053	0.06262	0.01047
b (loom)	1	86.83	14.46	13.15	58.42	6.125	0.062	0.01052
b (loom)	2	89.95	15.11	13.61	60.69	6.17	0.09744	0.01047
f (fla)	0	114.3	21.64	12.44	79.7	1.186	0.06658	0.01043
f (fla)	1	89.26	14.89	13.64	59.95	4.543	0.06069	0.0105
f (fla)	2	125.4	15.74	14.05	61.74	4.589	0.1241	0.01047

train data a vs b compares the per-rank training tensor dumps (--save-debug-train-data) of the two loom replays: 3 rollouts x 2 data-parallel ranks = 6 files, 5,382 tensors, all bit-identical. a vs b: MISSING refers to the separate grad-norm dump (--ci-save-grad-norm), which this tree does not write, so that channel has nothing to compare (same as in the previous comment); the grad norms are compared through the training logs instead (DET_COMPARE_LOGS, identical on all 3 steps). The replays ran with --load-debug-rollout-data ... --debug-train-only, i.e. no sampling; the third replay (f) trains the same rollouts with --linear-attention-backend fla for the same-data timing comparison above (its logprob diff / kl differ from loom's in the 4th digit, 0.01043 vs 0.01047 on step 0, as expected from a different kernel; its grad norms 0.0666 / 0.0607 / 0.1241 vs loom 0.0626 / 0.0620 / 0.0974).

Caveats

• step_time of step 0 includes the pre-training AIME eval (loom arm) and the first weight load; compare steady-state steps (>= 1) only.
• Rollouts are sampled, so the two arms see different trajectories; rewards, response lengths and the raw train_time of the sampled arms are therefore not directly comparable (the fla arm sampled about 11 % longer responses). The controlled train-time comparison is the same-data replay in the determinism section: the same three recorded rollouts are trained train-only with each backend.
• Host memory: CPU Adam holds about 113 GiB of optimizer state per training rank; the 8-GPU runs used a 1500 GiB memory cap as before.
• Loom backend evidence: the parsed-args dump shows linear_attention_backend loom, the loom short-conv backend selection is logged (short conv backend under loom: triton), and no fallback message appears (the only fallback match is the config print fallback_to_eager_attn False); the only Traceback blocks in the launcher logs are sglang's _freeze_gc_after_server_warmup readiness polls (connection refused before the HTTP ports open) and occur identically in the fla arm and in the previous comment's runs.
• The AIME numbers use 480 samples per eval and are noisy; the 10-step rewards are too few to read a trend.

🤖 Generated with Claude Code

[yyihuang]

Phase-2 kernels pushed: round-13 GDN chunk kernels (head 40ea59b)

The branch now carries the phase-2 kernel work: 17 commits on top of 9437f0b, all under miles_plugins/models/gdn_chunk_train (regenerated csrc/sm_100a, csrc/sm_103a, registry_sm_100a.json / registry_sm_103a.json with 25 stages each, ops.py, _jit.py) plus tests/fast-gpu/test_gdn_chunk_train_wiring.py. Nothing outside the plugin directory changed: the --linear-attention-backend loom selection, the weight layouts and the bit-determinism contract are as described above.

What the commits add

• Host side (2cc3aa0 … dde5249): one chunk table per micro-batch with the WY/state tensors kept on the autograd tape; segment-parallel state scans (pre-scan + stitch) so long sequences no longer serialize on one CTA per head; a calibrated segment planner that mirrors loom's plan_segments (prices seq_lens on the host, unit-tested); the wiring test creates its autograd leaves outside the profiled region.
• Kernel rounds 3 → 13 (one commit per round, each naming the loom commit it was exported from). Rounds up to 8 (loom ed1d12ef3ad) are the kernels the 2026-10-11 end-to-end run above used. Rounds 9–13 are new relative to that run; measured kernel-level, CUPTI GPU time of the GDN chunk kernels (forward, and forward+backward union), same-node interleaved A/B against the previous round's kernels, on B200 (sm_100a) and GB300 (sm_103a), best 16-row table entry:

round	loom commit	change	best-row gain vs previous round (sm_103a / sm_100a)
9	f94541bddab	WY pass tiles streamed through per-warp staging	fwd +5.65 % / +5.35 %
10	cf309778cc9	fwd_h split entry path scoped to the 64-row scan	fwd +7.49 % / +8.80 %, fwd+bwd +2.71 % / +3.25 %
11	1852a7e39eb	eight-warp WY pullback on a padded-pitch pool, stmatrix.x4 tiles, packed f32x2 epilogue	fwd+bwd +3.36 % / +3.32 %
13	9a9653333cc	forward WY on a padded-pitch pool with trace-time address folding and a line-aligned cp.async	fwd +5.1…+5.8 % / +4.9…+5.5 % on the 1 × 4096-token TP4 row (both launch orders); fwd+bwd +1.3 % / +1.1 % on that row

Rounds 12, 14 and 15 kept no candidate (every candidate stayed under the 2 % bar or was closed by the static census), and the optimization loop was closed by its stopping rule (two consecutive rounds without a ≥ 2 % best-row improvement).

Validation of the round-13 kernels (the ones on the branch now)

• Bit determinism: the determinism gate (12 shape rows × 100 repeats, forward + backward outputs hashed) is hash-equal on every row on both B200 and GB300.
• compute-sanitizer synccheck + racecheck on the plugin's launch path: 0 errors on both architectures (memcheck is not part of this gate).
• Kernel table regression check against the round-12 kernels: no cell more than 1 % slower (all 16 rows on GB300; the 10 targeted rows on B200), best row as in the table above.
• Export identity: re-rendering the plugin sources from the loom tree after it absorbed loom's main (today) reproduces all 25 stage symbols byte-identically on both architectures, so the committed csrc/ is exactly what the current loom tree generates.
• The plugin wiring / planner unit tests passed on both architectures in the same gates.

Not re-run, and still blocked

• The end-to-end Qwen3.8-27B numbers in the 2026-10-11 comment were measured with the round-8 kernels. The round-13 kernels on the branch were validated at kernel level only (above); the 27B recipe has not been re-run end to end with them. With the GDN chunk kernels at about 8 % of the train phase's device time on B200 (earlier measurement above), the kernel-level gains of rounds 9–13 translate into a small wall-time change at most — an expectation, not a measurement.
• The GB300 (TP2) end-to-end run of this recipe remains blocked by an sglang hang in the pinned radixark/miles:dev container; it was not resolved in this phase.

[yyihuang]

Follow-up to the push above: 0993335 adds the round-10 host-side dispatch that the kernel commits had left out.

The round-10 kernels export a split-plan instantiation fwd_h_r*_split of every forward scan stage (its prologue derives the segment's FP32 start state from the pre-scan maps and flags, so a segmented forward plan launches pre-scan + split scan and no forward stitch). The round-10 … round-13 commits shipped the csrc/ and registries but not the ops.py change that launches those stages, so the branch at 40ea59b still took the pre-scan + stitch + plain-scan path on segmented forward plans and never ran the split stages. 0993335 brings ops.py and the wiring test to the host code that every kernel-level gate since round 10 ran together with these kernels (wiring tests, determinism, A/B); the kernels, registries and everything else are unchanged.

I am re-running the Qwen3.8-27B end-to-end recipe (loom 10 steps + eval, fla 3 steps, same-data record/replay) on this head and will post the numbers here when it finishes.

[yyihuang]

Qwen3.8-27B dense GRPO re-run with the round-13 kernels (head 0993335)

Same recipe, node type (8 x B200, TP4), prepared checkpoint and datasets as the 2026-10-11 comment; the only change is the plugin on the branch head: round-13 csrc/ + registries and the round-10 host dispatch (ops.py). Before any arm ran, the plugin directory was overlaid from the branch head (50 files, 25 stages) and the wiring unit test passed (44 tests); the 25 stage extensions were JIT-built fresh in a private extension directory.

Summary

• Runs end to end with --linear-attention-backend loom: 10 rollout steps + AIME eval before and after, no OOM / NaN / training-side traceback, loom on all 8 ranks (no fallback). Step 0 rollout/raw_reward 0.6445, train/train_rollout_logprob_abs_diff 0.01049 (documented 0.64 / 0.010); AIME-2024 0.5417 (initial weights) → 0.6792 after the 10 steps (10-11: 0.5271 → 0.6417; 480 samples per eval, noisy).
• Steady-state wall time (steps 1–9): 293 s per step (rollout ≈ 198 s, train ≈ 82 s = ref log-probs 13.2 + actor log-probs 11.8 + actor train 55.5 + wait, weight sync ≈ 3.6 s); shares of the step: rollout 67 %, train 28 %, weight sync 1 %, the rest is the pre-/post-step wait. The GDN kernels' own device-time share inside the train phase was not profiled in this run (the recipe runs without a profiler); kernel-level numbers are in the push comment above (per-round gains of the kernel tables). 10-11 with the round-8 kernels: 295.5 / 200.1 / 81.0 (13.4 + 11.9 + 54.5) / 3.5.
• loom vs fla, sampled arms: step-time ratio 0.943 (293.0 vs 310.6 s), train-phase ratio 0.868 (81.7 vs 94.2 s); sampled response lengths differ again (the fla arm sampled about 12 % longer responses: 5471 vs 4882 tokens), and the length-normalized actor-train throughput is within 1 % (23,403 vs 23,228 tokens/s). The controlled comparison is the same-data replay below.
• Same-data replay (same three recorded rollouts, train-only, loom a / b vs fla f): actor train 64.3 vs 65.9 s (ratio 0.976), actor log-probs 12.6 vs 13.6, ref log-probs 15.9 vs 17.2, train_time 93.4 vs 97.3 s (ratio 0.960). Per step, train_time is lower with loom in every one of the 3 steps and in both loom replays (step 0: 105.0 / 104.5 vs 111.8 s; step 1: 85.4 / 85.5 vs 88.6 s; step 2: 89.1 / 91.0 vs 91.4 s), and so is the actor-train phase (75.5 / 74.5 vs 77.6; 57.7 / 57.1 vs 58.3; 59.9 / 61.0 vs 61.8); the ref log-probs pass is faster with loom on steps 0–1 and 0.1–0.3 s slower on step 2 (15.3 / 15.5 vs 15.2 s); the actor log-probs pass is faster on steps 0–1, and on step 2 one loom replay is 0.5 s faster and the other 0.1 s slower (13.4 / 14.0 vs 13.9 s). 10-11 (round-8 kernels): actor train 64.8 vs 67.1 s (0.966), train_time 97.7 vs 109.7 s.
• Bit determinism: the two loom replays produced bit-identical training tensors (5,382 tensors over 6 rollout/rank dump files, 0 differ); loss / pg_loss / entropy / grad-norm / kl identical in the logs.

loom arm (10 rollout steps)

step	raw_reward	response_lengths	truncated	train_rollout_logprob_abs_diff	train_rollout_kl	loss	grad_norm	rollout_time	ref_log_probs_time	log_probs_time	actor_train_time	train_time	update_weights_time	step_time	tokens_per_gpu_per_sec
0	0.6445	4589	0.3164	0.01049	0.0003811	-3.842e-09	0.1403	208.5	18.03	11.09	91.24	121.3	3.533	1859	704.5
1	0.4844	5412	0.5156	0.01056	0.0003917	4.424e-09	0.0589	193.1	14.76	13.14	60.78	89.81	3.414	307.1	897.2
2	0.5312	5502	0.4336	0.01046	0.0003849	-2.759e-08	0.1111	243.7	14.67	13.76	60.96	90.07	3.9	345.5	722.5
3	0.6406	4783	0.3477	0.009979	0.0003669	4.54e-09	0.06761	182.2	13.34	11.57	54.28	80.81	3.53	274.5	839.9
4	0.6289	4997	0.3672	0.01088	0.0003998	-1.397e-09	0.06901	186.3	13.24	12	56.43	83.5	3.633	280.7	858.5
5	0.5391	5175	0.4219	0.01026	0.0003802	-3.842e-09	0.1298	255.3	13.81	12.49	58.72	86.11	3.828	352.8	648.7
6	0.7148	4246	0.2852	0.009973	0.0003666	-7.451e-09	0.07294	152	11.73	10.28	49.09	72.82	3.576	236.3	893.7
7	0.625	5054	0.3672	0.01044	0.0003829	-4.191e-09	0.1004	229.6	13.4	12.17	57.11	84.5	3.532	332.7	704.5
8	0.7539	3846	0.2383	0.00997	0.0003671	-1.735e-08	0.112	145.6	10.74	9.31	46.39	67.03	3.6	223.9	845.1
9	0.6406	4924	0.332	0.01008	0.0003717	-5.588e-09	0.1265	190.8	13.17	11.86	55.4	80.97	3.68	283.4	825.8

Evaluation (AIME-2024, 30 problems x 16 samples; eval 0 = initial weights, eval 9 = after all 10 steps):

• eval 0: aime = 0.5417, aime/num_training_samples = 480, aime/truncated_ratio = 0.4562
• eval 9: aime = 0.6792, aime/num_training_samples = 480, aime/truncated_ratio = 0.3083

fla control arm (3 rollout steps)

step	raw_reward	response_lengths	truncated	train_rollout_logprob_abs_diff	train_rollout_kl	loss	grad_norm	rollout_time	ref_log_probs_time	log_probs_time	actor_train_time	train_time	update_weights_time	step_time	tokens_per_gpu_per_sec
0	0.6211	4706	0.3594	0.01052	0.0003805	-2.794e-09	0.07931	191.2	45.18	11.82	163.2	220.7	3.857	445.1	787.8
1	0.4922	5394	0.5039	0.01062	0.0003884	-4.657e-10	0.08599	190.1	16.6	14.76	61.52	93.44	3.992	302.8	908.1
2	0.5352	5547	0.4258	0.01049	0.0003841	-1.211e-08	0.09729	213.3	15.53	13.82	63.34	95	3.296	318.3	832.3

loom vs fla (same recipe, same machine)

metric	loom	fla	loom / fla
step time (mean over steps >= 1)	293	310.6	0.943
rollout time (mean over steps >= 1)	197.6	201.7	0.980
train time (ref + actor log-probs + actor train) (mean over steps >= 1)	81.74	94.22	0.868
ref log-probs time (mean over steps >= 1)	13.21	16.06	0.822
actor log-probs time (mean over steps >= 1)	11.84	14.29	0.829
actor train time (mean over steps >= 1)	55.46	62.43	0.888
weight sync time (mean over steps >= 1)	3.633	3.644	0.997
rollout tokens/GPU/s (mean over steps >= 1)	804	870.2	0.924
actor train tokens/s (length-normalized) (mean over steps >= 1)	2.34e+04	2.323e+04	1.008
mean response length (sampled) (mean over steps >= 1)	4882	5471	0.892
logprob abs diff (mean over steps >= 1)	0.01029	0.01055	0.975
train-rollout KL (mean over steps >= 1)	0.0003791	0.0003862	0.982

Bit determinism and same-data timing (record three rollouts with loom, replay them train-only twice with loom and once with fla)

### DET (06:04:14): record 3 rollouts with loom, replay them train-only with loom (a, b) and fla (f), compare; remaining=163min
### ARM det_record (06:04:14): num_rollout=3 --no-enable-eval extra='--linear-attention-backend loom --save-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt' output=<workspace>/... remaining=163min
ARM_DET_RECORD_EXIT=0 (24min) log=<workspace>/...
rollout dumps: 3 (expected 3)
### ARM det_replay_a (06:28:24): num_rollout=3 --no-enable-eval extra='--linear-attention-backend loom --load-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt --debug-train-only --ci-save-grad-norm <workspace>/det/grad_norm_a_{rollout_id}_{step_id}.pt --save-debug-train-data <workspace>/det/train_a_{rollout_id}_{rank}.pt' output=<workspace>/... remaining=139min
ARM_DET_REPLAY_A_EXIT=0 (12min) log=<workspace>/...
### ARM det_replay_b (06:40:34): num_rollout=3 --no-enable-eval extra='--linear-attention-backend loom --load-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt --debug-train-only --ci-save-grad-norm <workspace>/det/grad_norm_b_{rollout_id}_{step_id}.pt --save-debug-train-data <workspace>/det/train_b_{rollout_id}_{rank}.pt' output=<workspace>/... remaining=127min
ARM_DET_REPLAY_B_EXIT=0 (12min) log=<workspace>/...
### ARM det_replay_f (06:52:44): num_rollout=3 --no-enable-eval extra='--linear-attention-backend fla --load-debug-rollout-data <workspace>/det/rollout_{rollout_id}.pt --debug-train-only' output=<workspace>/... remaining=115min
ARM_DET_REPLAY_F_EXIT=0 (12min) log=<workspace>/...
### DET_COMPARE (07:04:53)
a vs b: MISSING
train data a vs b: BIT-IDENTICAL (5382 tensors equal, 0 differ over 6 rollout/rank files)
DET_COMPARE_EXIT=0 (2s)
### DET_COMPARE_LOGS (07:04:55)
train/loss: IDENTICAL (6 vs 6 values) ['1.862645149230957e-09', '1.862645149230957e-09', '-1.862645149230957e-09', '-1.862645149230957e-09'] vs ['1.862645149230957e-09', '1.862645149230957e-09', '-1.862645149230957e-09', '-1.862645149230957e-09']
train/pg_loss: IDENTICAL (6 vs 6 values) ['1.862645149230957e-09', '1.862645149230957e-09', '-1.862645149230957e-09', '-1.862645149230957e-09'] vs ['1.862645149230957e-09', '1.862645149230957e-09', '-1.862645149230957e-09', '-1.862645149230957e-09']
train/entropy_loss: IDENTICAL (6 vs 6 values) ['0.0', '0.0', '0.0', '0.0'] vs ['0.0', '0.0', '0.0', '0.0']
train/grad_norm: IDENTICAL (6 vs 6 values) ['0.07661361992359161', '0.07661361992359161', '0.11854755133390427', '0.11854755133390427'] vs ['0.07661361992359161', '0.07661361992359161', '0.11854755133390427', '0.11854755133390427']
train/kl_loss: IDENTICAL (6 vs 6 values) ['0.0', '0.0', '0.0004031201242469251', '0.0004031201242469251'] vs ['0.0', '0.0', '0.0004031201242469251', '0.0004031201242469251']
DET_COMPARE_LOGS_EXIT=0 (0s)
=== DONE Wed Oct 14 07:04:55 AM PDT 2026 rc=0 total=186min ===

Per-step timings of the three train-only replays of the same recorded rollouts (seconds; train_time = ref log-probs + actor log-probs + actor train + waits):

replay	step	train_time	ref log-probs	actor log-probs	actor train	train wait	grad_norm	logprob abs diff
a (loom)	0	105	17.2	11.24	75.47	1.345	0.07661	0.01061
a (loom)	1	85.44	14.24	12.76	57.75	5.097	0.1185	0.01053
a (loom)	2	89.09	15.3	13.37	59.88	5.439	0.08945	0.01051
b (loom)	0	104.5	18.19	11.25	74.54	1.527	0.07661	0.01061
b (loom)	1	85.45	14.87	12.92	57.13	5.211	0.1185	0.01053
b (loom)	2	90.96	15.5	13.95	60.98	5.585	0.08945	0.01051
f (fla)	0	111.8	20.92	12.67	77.56	1.339	0.07688	0.01058
f (fla)	1	88.63	15.43	14.36	58.26	5.203	0.1027	0.01056
f (fla)	2	91.44	15.16	13.87	61.75	4.474	0.09353	0.01054

Caveats

• step_time of step 0 includes the pre-training AIME eval (loom arm) and the first weight load; compare steady-state steps (>= 1) only.
• Rollouts are sampled, so the two arms see different trajectories; rewards, response lengths and the raw train_time of the sampled arms are not directly comparable. The controlled train-time comparison is the same-data replay.
• Loom backend evidence as before: linear_attention_backend loom in the parsed args, short conv backend under loom: triton logged by the training ranks (1 line + repeated 7x across cluster = 8 ranks), the only fallback match is the config print fallback_to_eager_attn False, the only nan match is the config print check_for_nan_in_loss_and_grad True, and the only Traceback blocks are one chained sglang _freeze_gc_after_server_warmup readiness poll (urllib3 connection refused before the HTTP port opens), identical in the fla arm and absent from the train-only replays.
• The GB300 (TP2) end-to-end run of this recipe remains blocked by the sglang hang in the pinned container (unchanged).
• AIME uses 480 samples per eval and is noisy; 10 steps are too few to read a reward trend.

🤖 Generated with Claude Code