[feat] move the HSTU arbitrary-mask path to NFUNC=5 - #670
Conversation
The fbgemm_gpu_hstu wheel bakes HSTU_ARBITRARY_NFUNC into its kernels and hard-checks func.size(-2) against it. It is being rebuilt with NFUNC=5 so OneRank can encode three visible column intervals (prefix + candidate + task token); the mask path exposes (NFUNC+1)/2 of them. SLA only needs two intervals, so build_sla_func_tensor pads the rest empty by repeating col_max1 -- every later (min, max) pair then satisfies max <= min, which is exactly the condition the kernel skips on. Target rows already relied on that behaviour. Every site hardcoded the literal 3; they now share HSTU_ARBITRARY_NFUNC, and the PyTorch reference decode ORs over (NFUNC+1)//2 intervals instead of two. OneRank's own func builder should use the same constant. The wheel pins and the install lines in the docs move together with the code: a 5-wide func against an fn3 wheel fails with "n_func must be equal to HSTU_ARBITRARY_NFUNC", so this cannot land before the fn5 wheels publish. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BertsxH7dKhSLHvTfgpvJj
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| # ``(NFUNC + 1) // 2`` visible column intervals per query row, and hard-checks | ||
| # ``func.size(-2) == NFUNC`` -- a mismatch is a TORCH_CHECK failure, so this | ||
| # constant and the installed wheel must move together. | ||
| HSTU_ARBITRARY_NFUNC = 5 |
There was a problem hiding this comment.
Non-blocking data request: the Verification section shows correctness parity and identical compile-time instantiation counts, but no runtime A/B of the fn5 vs fn3 wheel on an effectively-2-interval SLA workload. Two small costs land in the hot path: the kernel now evaluates (and skips) a third interval per query row in fwd and bwd, and the .contiguous() materialization of the stride-0 expand grows 5/3× (ABI-forced, no in-repo fix). Since both wheels share the same fece651b commit, a same-config SLA throughput + peak-memory comparison would be cheap and would confirm the skip is as inert at runtime as it is semantically.
There was a problem hiding this comment.
Measured, and the request was justified — the overhead is not negligible, which is what I had assumed before running it.
fwd+bwd of cutlass_hstu_mha on an SLA config, A10 (sm_86), bf16, 20 iterations after 5 warmup, fn3 wheel paired with master's 3-wide builder and fn5 with this branch's 5-wide one (both fece651b):
| config | fn3 | fn5 | delta |
|---|---|---|---|
| B=32 S=1024 H=4 D=128 K1=256 K2=32 | 2.393 ms | 2.561 ms | +7.0% |
| B=16 S=2048 H=8 D=128 K1=256 K2=32 | 5.923 ms | 6.479 ms | +9.4% |
Run-to-run spread was <= 0.07 ms across 3-4 runs per leg, so both deltas are outside noise. Peak allocated memory moved 329.9 -> 331.1 MiB on config A, i.e. the +1.0 MiB func buffer plus change; the buffer itself goes 1.50 -> 2.50 MiB and 3.00 -> 5.00 MiB.
The .contiguous() growth you flagged is not where the time goes: 2 MiB extra at ~600 GB/s is ~3 us against a 556 us delta. It is kernel-side — 50% more warpReduce passes in the prologue (the gMinFunc/gMaxFunc reductions run over kNFunc/2 + 1 rows) plus a third interval evaluated in the per-block masking loop. So SLA workloads pay 7-9% on the attention op for an interval that is always empty; OneRank, which uses the third interval, gets the value. Scope worth stating: that is the attention op, not an end-to-end training step, and these are Ampere-path numbers on one GPU.
Separately, while chasing the memory half of this: the nheads dimension of func is never traversed by any backend. Hopper hardcodes the head stride to zero — make_stride(/*mainloop_params.func_head_stride*/Int<0>{}, ...) at hstu_fwd_kernel.h:224 and hstu_bwd_kernel.h:253 — and Ampere gives the head extent Int<1> with params.h commented out (hstu_fwd.h:157, hstu_bwd.h:153) and then indexes it at 0. Neither backend has a CHECK_CONTIGUOUS on func; only dtype and device are checked. So the .contiguous() on the stride-0 expand(nheads, NFUNC, total_q) materializes nheads identical copies of which the kernel reads one — at nheads=8, total_q=1e6 that is 160 MB where 20 MB would do, and the waste predates this PR (96 MB at fn3).
Not folding that in here: it changes a shipped path, needs GPU verification, and those Int<0> hardcodes are commented-out code a future fork bump could restore, at which point a (1, NFUNC, total_q) tensor would silently feed wrong rows to heads > 0. Filed as its own follow-up with a guard.
| and attn_func.shape[1] == HSTU_ARBITRARY_NFUNC | ||
| and attn_func.shape[2] == q.shape[0], | ||
| "attn_func must have shape (nheads, 3, total_q)", | ||
| f"attn_func must have shape (nheads, {HSTU_ARBITRARY_NFUNC}, total_q); " |
There was a problem hiding this comment.
This message helps for hand-built tensors, but it can't fire for the realistic skew the docs set up: build_sla_func_tensor always matches the repo constant, so a stale image with the pre-installed fn3 wheel + this code passes every Python-side assert and fails later at the kernel's TORCH_CHECK — potentially minutes into a job, after AOT compile. Since the new wheels carry a machine-readable fn<N> tag, a one-time importlib.metadata.version("fbgemm_gpu_hstu") check (warning or error) at first use of the attn_func path would turn that into an actionable "installed wheel is fn3, code requires fn5". Raised as a question, not a blocker — AGENTS.md discourages speculative defenses, and the failure is already clean (exception, not memory unsafety).
There was a problem hiding this comment.
The gap is real — you're right that build_sla_func_tensor always matches the repo constant, so a stale image with the pre-installed fn3 wheel sails past every Python assert and dies at the kernel TORCH_CHECK on the first attention forward. Not adding the importlib.metadata check here, though, for two reasons.
It cuts against the repo's established shape. There is no importlib.metadata anywhere in tzrec/; the only two runtime version checks (env_util.enable_tma() and its test-util twin) use packaging.version on a __version__ attribute and warn rather than raise; and dynamicemb_util.py:207 ships a torchrec compat shim for a known behaviour change without reading torchrec's version at all. Against AGENTS.md:49 ("add no speculative features, abstractions, configurability, fallbacks, or defenses") a new dependency-introspection pattern for a failure that is already a clean exception is hard to justify.
More to the point, parsing fn5 out of a version string is a proxy for a value the wheel should simply expose. HSTU_ARBITRARY_NFUNC is a preprocessor macro that reaches the TORCH_CHECK and the kernel template but never Python — the four registered ops (hstu_varlen_{fwd,bwd}_{80,90}) neither take nor return it, and nothing in the hstu package namespace reports it. The tag only exists because we added it to BUILD_VERSION, so it is absent on older wheels (0.1.0+20260823.fece651b.cu130) and on upstream builds (0.1.0+cu12.6) — a best-effort signal, not a guarantee.
The fix belongs in the fork: a three-line op next to the existing symint.patch,
int64_t hstu_arbitrary_nfunc() { return HSTU_ARBITRARY_NFUNC; }
m.def("hstu_arbitrary_nfunc() -> int");registered once (the ampere TORCH_LIBRARY_FRAGMENT; registering in all three would collide). The macro is already passed to both cxx and nvcc, so it is in scope, and it usefully returns 0 under HSTU_DISABLE_ARBITRARY=TRUE. That reads the real value instead of a string, needs no guessing about absent tags, and is upstreamable. It costs a wheel rebuild, so it will land at the next fork bump rather than in this PR — tracked for that bump.
| torch.testing.assert_close(col_min, col_max1) | ||
| torch.testing.assert_close(col_max, col_max1) | ||
| self.assertTrue(bool(torch.all(col_max <= col_min)), f"interval {i}") |
There was a problem hiding this comment.
The assertTrue(all(col_max <= col_min)) is dead code: the two assert_close calls above already pin both rows to col_max1, so if they pass the inequality is trivially true, and if they fail execution never reaches line 297. The assertion documenting the actual kernel contract (max <= min ⇒ skipped) can therefore never fire independently, while the equality asserts over-pin the padding scheme — a future equally-valid empty encoding (e.g. min == max == 0) would fail this test despite identical kernel behavior. Consider keeping just the contract assertion, or dropping the redundant assertTrue.
Related, optional: on wheel-free environments the only decode coverage is test_sla_matches_fixed_causal_when_k1_spans_full_window, whose full-window parameters make it structurally blind to a pad-interval leak (any row-index bug in the generalized loop still yields intervals that are empty or subsets of [0, q+1)). A small CPU test comparing pytorch_hstu_mha(attn_func=...) against a hand-built expected mask with non-degenerate k1/k2 would cover rows 3–4 semantically; defensible to accept as-is since the GPU parity test runs in the declared CI lanes.
There was a problem hiding this comment.
Both points were right; fixed in 4d7580a.
The assertTrue is gone as dead code, and the two assert_close calls went with it — you're right that they pinned the padding scheme rather than the kernel contract, and min == max == 0 would have been an equally valid encoding that failed the test. What's left is the width check plus the emptiness contract, which is now the live assertion. Confirmed it can actually fail: padding with [col_max1, col_max1 + 1] trips it with interval 2 is not empty.
On the CPU-coverage note — this turned out to be the substantive one, and it is now covered by test_decode_covers_all_nfunc_intervals. Your reasoning was exactly right: with sla_k1 = N, sla_k2 = 0 the rows are [0, 0, q+1, q+1, q+1], so intervals 0 and 2 are empty and any misread of rows 3-4 stays invisible.
Rather than hand-build an expected mask, the new test partitions the causal row [0, q+1) into (NFUNC + 1) // 2 adjacent non-empty chunks — [max0, min0, max1, min1, max2] = [s1, s1, s2, s2, q+1] with s1 = min(2, q+1), s2 = min(5, q+1). The union is again exactly causal, so it compares against pytorch_hstu_mha(causal=True) the same way its neighbour does, but every row now carries a live boundary.
Verified by mutation. Changing the decode loop to range(1, (n_func + 1) // 2 - 1) — silently dropping the last interval — gives:
test_decode_covers_all_nfunc_intervals FAILED
test_sla_matches_fixed_causal_when_k1_spans_full_window PASSED
which is precisely the blindness you described. It goes through the public pytorch_hstu_mha, not _decode_attn_func_to_mask, since no test in this repo imports a private helper, and it carries no GPU skip so it runs on the unscoped CPU lane.
| ``(nheads, 3, total_q)``, int32 — selects the NFUNC mask | ||
| path. Supported on ``Kernel.CUTLASS`` and ``Kernel.PYTORCH``; | ||
| rejected on ``Kernel.TRITON``. | ||
| ``(nheads, HSTU_ARBITRARY_NFUNC, total_q)``, int32 — selects the |
There was a problem hiding this comment.
Minor: hstu_mha is the main public entry, but this module neither imports nor defines HSTU_ARBITRARY_NFUNC, so a reader hits the symbol here with no way to resolve it to "5" (or to the fn5 wheel requirement) without grepping. The backend modules that mention it do import it. A parenthetical like "(defined in tzrec.ops.hstu_attention_utils; currently 5, must match the installed wheel's fn<N> tag)" would make the docstring self-contained.
There was a problem hiding this comment.
Fixed in 4d7580a — the docstring now names tzrec.ops.hstu_attention_utils as the definition site, gives the current value, and states the fn<N> wheel correspondence. Docstring only; no import added for a docstring's sake.
Code review (5-area subagent review: quality, performance, test coverage, docs, security)Overall: LGTM. The change is surgical and internally consistent. Independently verified:
Four non-blocking notes posted inline:
🤖 Generated with Claude Code |
Test the kernel contract, not the padding scheme. The pad-interval test asserted `col_min == col_max1` and `col_max == col_max1` and then, behind those, `col_max <= col_min` -- which could never fail independently, since the equalities imply it. The equalities also over-pinned the encoding: a future `min == max == 0` padding is equally valid to the kernel and would have failed the test. Only the emptiness contract is asserted now, and it is live: padding with a non-empty interval trips it. Cover every interval the compiled NFUNC exposes. `build_sla_func_tensor` fills only the first two, so the rows past them were never exercised semantically -- and the one CPU test that decodes a func tensor uses `sla_k1 = N, sla_k2 = 0`, which makes intervals 0 and 2 empty anyway. A decode loop that silently drops the last interval still passes it. The new test partitions the causal row into (NFUNC + 1) // 2 adjacent non-empty chunks whose union is again plain causal, so every row carries a live boundary; it fails on exactly that mutation while the existing test passes. No GPU skip, so it runs on the CPU lane. Point the `hstu_mha` docstring at where HSTU_ARBITRARY_NFUNC is defined -- the module neither imports nor defines it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BertsxH7dKhSLHvTfgpvJj
…aba#670) Review [2a]: now that alibaba#670 landed HSTU_ARBITRARY_NFUNC = 5, candidate replicas are retired and each candidate expands into exactly K + 1 tokens, as the paper writes it. - build_onerank_func_tensor: 5-wide three-interval encoding. t_k's visible set [0, H_b) u {e^C_i} u {t_k} splits naturally into the (NFUNC + 1) // 2 = 3 column intervals because the mutually invisible t_1..t_{k-1} sit between e^C_i and t_k -- the very reason replicas existed under NFUNC=3. - OneRankTokenizer: group_size is K + 1; forward concatenates the candidate with the task tokens (torch.cat, no replica expand). - Inflation factors follow across tests and docs: default max_seq_len 148 -> 132, runtime-rejection case 26 -> 18, docs example 2328 -> 2208. - group_size validation drops the even-parity rule (odd sizes are legal now); the floor is 2. - Tests rewritten against the paper layout: reference mask, slot indices, stamp positions; the replica-gradient test is replaced by one covering the single candidate row plus the per-group task-token gradient aggregation. - CUTLASS equivalence tests now require the fn5 wheel (CI h20 lane); locally verified through the PyTorch reference kernel (_decode_attn_func_to_mask) instead.
3 participants
Why
OneRank (#661) encodes task-private attention masks as per-row visible column intervals on the HSTU arbitrary-mask path, and needs three of them (prefix + candidate + task token) so it can drop its candidate replicas. The kernel exposes
(NFUNC + 1) / 2intervals, so thefbgemm_gpu_hstuwheel has been rebuilt withHSTU_ARBITRARY_NFUNC=5(2 intervals -> 3).NFUNCis a compile-time constant of the wheel — the host side hard-checksfunc.size(-2)against it in all three backends (ampere / hopper / blackwell_rtx) and raisesn_func must be equal to HSTU_ARBITRARY_NFUNC. Nothing reports it at runtime, so the new wheels carry it in their version as anfn5segment.What changed
tzrec/ops/hstu_attention_utils.py: newHSTU_ARBITRARY_NFUNC = 5constant. Every site previously hardcoded the literal3; they now share this one.build_sla_func_tensoremits[col_max0, col_min0, col_max1]plus pad rows. SLA only needs two intervals, so the pad repeatscol_max1— every later(min, max)pair then satisfiesmax <= min, which is exactly the condition the kernel skips on (hstu_fwd.h:if (f_max <= f_min) { continue; }). Target rows already relied on that behaviour._decode_attn_func_to_mask(PyTorch reference) ORs over(NFUNC + 1) // 2intervals instead of a hardcoded two, matching the kernel's interleaved[max0, min0, max1, min1, ...]row layout.requirements/cu{126,129,130}.txtand the install lines indlrm_hstu.md/ultra_hstu.mdmove to0.1.0+20260914.fece651b.fn5.${DEVICE}.The pins and the code must land together: a 5-wide
funcagainst an fn3 wheel (or vice versa) fails theTORCH_CHECK. The fn5 wheels are published for cu126/cu129/cu130 × py310/311/312.Build-side change: alibaba/TorchEasyRec-addons MR #29992875.
Verification
test_sla_attn_cutlasspasses on an A10 with the fn5 wheel installed — CUTLASS output matches the PyTorch reference with the 5-widefunc, which is what proves the padded third interval is inert.test_sla_matches_fixed_causal_when_k1_spans_full_window(CPU) still reproduces plain causal attention, exercising the generalized decode.test_width_matches_nfunc_and_extra_intervals_are_emptypins the tensor width and the emptiness of the pad intervals..sogrew 871 MB -> 877 MB.🤖 Generated with Claude Code
https://claude.ai/code/session_01BertsxH7dKhSLHvTfgpvJj