Repository navigation
Qwen3.6-35B-A3B on CUDA: a 1M-token window with static YaRN rope - #1675
Merged
Merged
Conversation
This was referenced Oct 30, 2026
laudney
force-pushed
the
pr-q36-yarn
branch
3 times, most recently
from
October 31, 2026 13:40
93fb4cf to
a9ba340
Compare
Owner
|
Thanks @laudney, both rebases are clean and prepared here on current main with your commits as they are. They're in our gate now, and #1675 (which carries #1674) is queued on a GB10 to check that the unset path stays exact and level on Qwen3.6-35B-A3B. We don't yet have |
Owner
|
Device results, @laudney: on a GB10, Qwen3.6-35B-A3B MLX 4-bit with its MTP head (the default |
On CUDA the Qwen3.6 MoE target wired --drafter to its MTP layer alone, though the engine's DFlash2 drafter already reads the SGLang recipe's incoai/Qwen3.6-35B-A3B-DFlash2. One openDrafter now serves both families: a directory whose config.json has dflash_config opens DFlash2, as every directory does under a dense target, and any other opens the MoE target's MTP layer. Both openers return one Opened struct whose prepare is optional, which cuda_drafted.zig now accepts. On a GB10 one build drafts with either; every drafted reply equals its plain run, and DFlash2 decodes 130.3 tok/s at one stream on a 4.2K-token prompt against MTP's 98.6.
The Qwen CUDA path needs transformers' static YaRN frequencies for its 64 rotary dims, which only flashnext computes today, in a file its Metal accuracy receipts hash by bytes. core.rope_yarn copies the formula so that file stays as it is; a host test holds the copy to transformers 5.19.0's own fp32 words for Qwen3.6's config, and a Flash Next test holds the two copies equal over six configurations.
Several Qwen3.5/3.6 config readers ignored rope scaling, so a YaRN config such as SGLang's staged 1M copy of Qwen3.6 would load and run with plain rope: HIP's and loader.zig's config.geometry, the dense Metal reader, the CUDA geometry's rope_scaling block, and both DFlash2 parsers. Each now refuses a rope_type or type other than default and a non-null rope_scaling by name, and config.geometry gets the trained-window limit parse already had; published configs load as before, with a test for each refusal.
The Qwen target's rope and the MTP layer's turn each pair by theta's own frequency, so a YaRN window past 262,144 positions has nowhere to go. tf_q35_rope, tf_q35_rope8 and tf_q36m_norm_rope now take a 33-float table (the attention factor, then 32 frequencies); a null table keeps today's expression. The engine holds the table's device address, the forward passes it, and the MTP's serial and batched layers share one launch helper. A new FP64 contract (rope-yarn) covers text and image rows to position 1,048,575 on bf16 and e4m3 caches and the MTP kernel's q and k, and the records follow the kernels' new register use.
SGLang serves Qwen3.6 at a 1M window from a staged config with YaRN rope; TensorFold refused such a config and had no way to stretch a downloaded one. cuda_rope.resolve now reads a YaRN block (in rope_parameters or the older rope_scaling) within transformers' default formula, or applies --rope-yarn FACTOR over the config's trained window, and the CUDA server checks --context against that ceiling. The engine uploads the table before any forward and again for the flag, dropping graphs captured with the old one. Every other CUDA family, HIP's open and follow, Metal and XPU refuse the flag by name; two --parallel helpers move to cli_parallel.zig so cli.zig keeps its 600-line limit.
Drafted replies equal their plain runs whatever the drafter's rope, since the target verifies every draft, so that gate cannot show the MTP layer missing its YaRN table. The drafter's GPU check now sets a table on the engine, absorbs three rows at position 300,000 and holds each ring key to the float64 YaRN rope of the key row its layer computed, then runs the batched-against-serial check with the table; it passes on a GB10 within 0.29 of its bound.
Commit 8317b03a merged the plain and YaRN rope into one expression, and nvcc, contracting with fmad on, then fused plain rope's pair products differently: three of six served reply hashes changed with YaRN off. The target's rope and the MTP layer's now branch once per launch into a plain instantiation that is the old code, and a YaRN one; on a GB10 every plain reply hash equals main's again, and the YaRN contracts and the MTP checks still pass. The records follow: each instantiation keeps its own shared row, 512 bytes more a block.
The recipe now shows how to serve nvidia/Qwen3.6-35B-A3B-NVFP4 at 1,048,576 tokens, by --rope-yarn 4 or a staged YaRN config, what it costs a stream, what YaRN changes, and the GB10 measurements: exact drafted, concurrent and resumed replies, a passphrase found at 300,000 and near 1M tokens, and the window's boundary; the CHANGELOG has the entry.
The YaRN table changed tf_q35_rope and tf_q35_rope8, and their 13.0 records were measured again, but nvcc 13.2's and 13.3's override entries still carried the old kernels' usage, which the registry prefers for images those compilers build. Both entries go, so those compilers' images fall back to the records' pending state until a 13.2 or 13.3 receipt of the new kernels is measured.
Sliding Weights' backward attention ropes with theta's plain frequencies and no attention factor, so under --rope-yarn or a YaRN config its gradients would not match the forward pass they train. The learner's shared admission, cuda_learn_shape.from, which both the server's --slide and tf-qwen35-learn reach after the engine knows its rope, now refuses a YaRN rope by name instead of learning with the wrong one.
Owner
|
Thank you @laudney. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1656 (3 of 3). Builds on #1674 (DFlash2): review only the last nine commits, from
feat(core): transformers' YaRN table for rope users outside flashnext; the first commit is #1674's and drops out once it merges. It needs #1674 because its draft-config rope check extends code that one adds. Rebased ontomainafter 1.2.0: the qwen27 rope refusal test now sits inconfig_test.zig,--rope-yarnincli_gpu_flags.zig(the earliercli_parallel.zigsplit is no longer needed), and the YaRN learner refusal sits beside the KV-share messages.Why
The checkpoint was trained on 262,144 positions, and SGLang's Spark recipe serves NVIDIA's
nvidia/Qwen3.6-35B-A3B-NVFP4at 1,048,576 with static YaRN (factor 4). Serving the same published weights at the same window with the same rope is what makes a long-context comparison like for like, and long-context users can run the checkpoint as downloaded (--rope-yarn 4) or SGLang's staged 1M copy as is, instead of a TensorFold-specific copy.What
core/rope_yarn.zig): transformers' static YaRN table (_compute_yarn_parameters, beta 32 and 1, truncated ends, torch's fp32) for rope users outside Flash Next; Flash Next keeps its own copy, which its Metal receipts hash.--slide(the learner) under YaRN.--rope-yarn 4on the checkpoint as downloaded, or a config whose rope carries a YaRN block, raises the window to 1,048,576 (--context 0serves it whole). A block and a flag must name the same factor. Without either, every byte is as before.Gate (one DGX Spark: GB10, sm_121, CUDA 13.0; checkpoint
1355db6a)The run on record is the recipe's, with the changelog's numbers:
--parallel 2, 22.2 GiB a stream, so a GB10 admits two): drafted replies equal their"draft": falseruns (6 of 6), concurrent ones their solo runs, a resumed second turn its fresh run. A passphrase 5% into a 300,000-token prompt comes back after 173 s of prefill (1,740 tok/s). A 1,048,512-token prompt prefills in 1,546 s (680 tok/s), returns its passphrase, and decodes 64 tokens to the window's last position at 31 tok/s; one token more is refused with the server's context-window error.--rope-yarn 4on the download and SGLang's staged 1M config give the same reply hashes (6 of 6).tf-cuda-test rope-yarn(YaRN rope intf_q35_rope,tf_q35_rope8andtf_q36m_norm_ropeagainst float64, to 1M) and the MTP check pass.zig build testandzig build test-cuda-host(89/89) pass;zig fmt --checkandgit diff --checkare clean;lean_check.pyreports 426 problems, the same asmain.A re-run on this stack passed every check above again (the 300,000-token needle included), without repeating the 26-minute 1M prefill.
Overlap
git merge-treeagainst every open pull request: most shared files are adjacent additions (a field, an error, a command or a test beside another PR's). The one real overlap is #1632: both change hownative/cuda.zigfinds a family's context ceiling (this PR'scontextCeilingfor YaRN, #1632'strained_context); merged, it checkscontextCeiling, thentrained_context, thenmodelContext.