Qwen3.5-35B-A3B-FP8: GDN decode lossy fast path + fused MRoPE QK

[zovonoir]

Qwen3.5-35B-A3B-FP8: GDN decode lossy fast path + fused MRoPE QK

Summary

This PR ships two complementary decode-time optimizations for Qwen3.5-35B-A3B-FP8. The MRoPE QK fused path was originally in #888 and has been merged here per review feedback that the two changes should land together.

1. GDN decode recurrent fast path

A decode-only fast path for ATOM GDN recurrent attention. The optimized path is enabled only for the common non-spec decode case:

• no prefill
• no speculative sequence mask
• one decode token per active request
• 1D ssm_state_indices
• a/b shaped as [num_decodes, local_num_v_heads]

For unsupported layouts, the existing fused_gdn_gating + fused_recurrent_gated_delta_rule fallback is preserved.

The fast path fuses GDN gating computation with the recurrent state update and avoids launching the separate gating kernel. The implementation is intentionally named gdn_decode_update_lossy_fast because it is an approximate decode path: it keeps the high-performance kernel structure and may introduce small bf16-scale numerical differences in the recurrent state update. End-to-end GSM8K validation shows accuracy remains on par with main.

The path is guarded by ATOM_ENABLE_GDN_DECODE_LOSSY_FAST and defaults to disabled. Set ATOM_ENABLE_GDN_DECODE_LOSSY_FAST=1 to enable the approximate fast path for A/B testing or performance runs.

PAD_SLOT_ID (-1) handling: the kernel writes zeros to out and skips state load/store for idle/padded slots, so downstream consumers of the full out buffer never see uninitialized memory.

2. Fused MRoPE Q/K Triton path

A specialized Triton MRoPE Q/K fusion for the Qwen3.5 hot path, wired into Qwen3NextAttention after Q/K norm and before attention. Falls back to the generic rotary embedding path for unsupported shapes (head_size != 256, non-MRoPE positions, etc.).

The dispatcher try_mrope_qk_fused is decorated with @torch.compiler.disable so the Python-level shape branches (e.g. positions.shape[0] != 3, q.shape[1] != num_q_heads * head_size) never reach Dynamo. Otherwise Dynamo specializes the symbolic dims to constants, conflicting with dims the runner marked dynamic and tripping ConstraintViolationError under MMStar dynamic-shape compile. This matches the convention used by chunk_gated_delta_rule in atom/model_ops/fla_ops/chunk.py.

Notes on the external SGLANG_MAMBA_SSM_DTYPE=bfloat16 setting

The opt+env measurements below also used the external SGLang runtime setting SGLANG_MAMBA_SSM_DTYPE=bfloat16 to store the Mamba/GDN SSM state in bf16. That setting is not introduced by this PR; it is included only as an additional benchmark configuration. Because the decode kernel is strongly memory-bandwidth bound, bf16 SSM state provides additional throughput improvement, with a larger accuracy tradeoff in the measured GSM8K run.

Performance

Benchmark config (MI308X):

• Model: Qwen3.5-35B-A3B-FP8
• CONC: 224, ISL: 4094, OSL: 2048, Prompts: 448
• TP/EP: 1/1
• Runtime: SGLang FP8

Configurations:

• main: baseline.
• opt (GDN only): opt-qwen35b with ATOM_ENABLE_GDN_DECODE_LOSSY_FAST=1, MRoPE fused path disabled.
• opt (GDN+MRoPE): this PR with ATOM_ENABLE_GDN_DECODE_LOSSY_FAST=1, fused MRoPE active.
• opt+env: opt (GDN only) plus external SGLANG_MAMBA_SSM_DTYPE=bfloat16.

Metric	main	opt (GDN only)	opt (GDN+MRoPE)	opt+env	GDN+MRoPE vs main
Total token throughput (tok/s)	7466.90	7870.39	8004.41	8103.82	+7.20%
Input token throughput (tok/s)	4977.12	5246.07	5335.41	5401.67	+7.20%
Output token throughput (tok/s)	2489.78	2624.32	2669.01	2702.15	+7.20%
Request throughput (req/s)	1.22	1.28	—	1.32	—
Benchmark duration (s)	368.51	349.62	—	339.55	—
Mean E2E (ms)	176401	167139	164893	162062	-6.52%
Mean TTFT (ms)	17782	17852	17771	17824	-0.06%
Mean TPOT (ms)	77.44	72.93	71.87	70.46	-7.19%
P99 TPOT (ms)	85.47	81.00	79.83	78.46	-6.60%

For reference, a fake-return experiment that bypassed fused_recurrent_gated_delta_rule entirely showed an approximate upper bound of QPS 1.4882, TPOT 61 ms, Total throughput 9140 tok/s — confirming the GDN recurrent update is a meaningful decode-time bottleneck and this PR captures a conservative portion of the opportunity.

Accuracy

GSM8K, 5-shot:

Version	flexible-extract	strict-match	flexible delta vs main	strict delta vs main
main	0.8946 ± 0.0085	0.9052 ± 0.0081	—	—
opt (GDN only)	0.8984 ± 0.0083	0.9090 ± 0.0079	+0.0038	+0.0038
opt (GDN+MRoPE)	0.895 ± 0.0097	0.903 ± 0.0094	+0.0004	-0.0022
opt+env (GDN + bf16 SSM)	0.8825	0.8931	-0.0121	-0.0121

The combined GDN+MRoPE accuracy remains on par with main. opt+env provides the highest throughput but with a measurable GSM8K drop in this run.

Cross-model benchmarks on Qwen3 dense (e.g. 27B) and other MoE models (e.g. 397B) are in progress and will be appended once available.

Validation

Additional standalone kernel checks were run for the target shape:

• N=224, H=16, HV=32, K=128, V=128, dtype=bfloat16

The optimized fast path may produce small bf16-scale differences versus the reference recurrent path. End-to-end GSM8K validation and decode benchmark results above were measured with the current gdn_decode_update_lossy_fast + fused MRoPE implementation.

The existing fallbacks remain in place for prefill, speculative decode, mixed layouts, unsupported metadata shapes, and unsupported MRoPE shapes.

Files

• atom/model_ops/attention_gdn.py — wire gdn_decode_update_lossy_fast into the GDN attention forward path; restore core_attn_out tail-zeroing for CUDA-graph padding safety.
• atom/model_ops/fla_ops/__init__.py — export gdn_decode_update_lossy_fast.
• atom/model_ops/fla_ops/fused_recurrent.py — add gdn_decode_update_lossy_fast + Triton kernel with PAD_SLOT_ID handling, contiguous-state validation, and HV % H checks.
• atom/model_ops/triton_mrope.py — new specialized Qwen3.5 MRoPE Q/K Triton kernels (tiled + per-token) with @torch.compiler.disable-decorated dispatcher.
• atom/models/qwen3_next.py — wire try_mrope_qk_fused into Qwen3NextAttention.
• atom/utils/envs.py — add ATOM_ENABLE_GDN_DECODE_LOSSY_FAST env var.
• tests/test_envs.py — include new env var in _ATOM_ENV_VARS cleanup list and add override assertion.

[Copilot]

Pull request overview

This PR adds a dedicated Triton “fast decode” update kernel for Gated Delta Net (GDN) and wires it into the Qwen3.5 decode path to improve decode throughput by avoiding materializing gates and using a streamlined state update.

Changes:

• Introduces gdn_decode_update_fast (Triton kernel + Python wrapper) for decode-time state update and output computation.
• Updates atom/model_ops/attention_gdn.py to conditionally use the fast decode kernel for non-spec decode cases and to skip gate materialization when not needed.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
atom/model_ops/fla_ops/gdn_decode_fast.py	Adds a new Triton kernel + wrapper intended to accelerate GDN decode updates.
atom/model_ops/attention_gdn.py	Integrates the new fast decode path and adjusts gating computation to be conditional.

Comments suppressed due to low confidence (1)

atom/model_ops/fla_ops/gdn_decode_fast.py:170

• The kernel launch passes threshold=threshold, but _gdn_decode_update_kernel does not have a threshold parameter (it currently declares threshould). This will fail at runtime when launching the Triton kernel; make the argument name consistent between the kernel signature and the launch call.

        bv,
        n_value_heads // n_heads,
        num_warps=4,
        num_stages=1,
        threshold=threshold
    )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (5)

atom/model_ops/attention_gdn.py:30

• This try/except currently re-imports and overwrites gdn_decode_update_fast, but due to the earlier unconditional import it is effectively dead code (and can’t protect startup). After switching to a guarded import, consider using a single import style (either try: from ... import ... or importlib.import_module) to avoid shadowing/duplication.

try:
    gdn_decode_update_fast = importlib.import_module(
        "atom.model_ops.fla_ops.gdn_decode_fast"
    ).gdn_decode_update_fast
except ImportError:
    gdn_decode_update_fast = None

atom/model_ops/fla_ops/gdn_decode_fast.py:167

• HEADS_PER_V is computed as n_value_heads // n_heads and then used as a divisor in the kernel (i_h = i_hv // HEADS_PER_V). If n_value_heads < n_heads or the ratio isn’t an integer mapping, this will either divide by zero or silently mis-map heads. Add explicit validation (e.g., n_value_heads % n_heads == 0 and n_value_heads >= n_heads) before launching the kernel.

        bk,
        bv,
        n_value_heads // n_heads,
        num_warps=4,

atom/model_ops/fla_ops/gdn_decode_fast.py:141

• cu_seqlens is required (ValueError if None) but is not used anywhere in the fast path. Either remove this requirement or use cu_seqlens to validate that the input really matches the supported decode-only shape (e.g., one token per sequence / cu_seqlens[-1] == n_tokens) so incorrect calls fail early for the right reason.

    if cu_seqlens is None:
        raise ValueError("gdn_decode_update_fast requires cu_seqlens for decode")
    if num_accepted_tokens is not None:
        raise ValueError("gdn_decode_update_fast does not support spec decoding")

atom/model_ops/fla_ops/gdn_decode_fast.py:156

• The wrapper doesn’t enforce contiguity/strides for q/k/v but the kernel uses implicit contiguous flattening (no strides passed). To match existing patterns (e.g., fused_recurrent_gated_delta_rule), consider calling .contiguous() on q, k, and v (and potentially out) or passing explicit strides to the kernel.

        a,
        dt_bias,
        q,
        k,
        v,
        b,
        out,

atom/model_ops/fla_ops/gdn_decode_fast.py:127

• gdn_decode_update_fast implicitly assumes the flattened varlen layout used elsewhere (batch dimension == 1). Right now it takes q.shape[0]/v.shape[0] as _ and then uses squeeze(0) plus linear indexing that will be incorrect if q/v have batch > 1. Add an explicit check (e.g., q.shape[0] == v.shape[0] == 1) and raise a clear error if not.

    _, n_tokens, n_heads, head_k = q.shape
    _, _, n_value_heads, head_v = v.shape
    out = torch.empty_like(v).squeeze(0) if o is None else o.squeeze(0)

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 5 comments.

[zovonoir]

Accuracy dropped confirmed at latest docker image, still debugging.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.

[zovonoir]

1. Please rebase main and then check the accuracy CI test.
2. It seems there is a 5% performance gain. Could you provide more details on how the above optimization methods perform on Qwen3 dense models (e.g., 27B) and other MoE models (e.g., 397B)?

testing in progress

[zovonoir]

@wanzhenchn
Here are tesing results
main : main branch without this patch
opt: this patch
bf16: prepend "SGLANG_MAMBA_SSM_DTYPE=bfloat16"
fp32: delete "SGLANG_MAMBA_SSM_DTYPE=bfloat16"

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

[zovonoir]

@wanzhenchn thanks for the review. Status on your two requests:

1. Rebase main + accuracy CI: done. main has been merged into the branch (commit 7d7f4fa), bringing in 20 commits including [fix](mtp): fix qwen3.5 mtp accept rate zero #966 (qwen3.5 mtp accept-rate fix) and scheduler: fix dead HBM prefix cache under deferred-output mode (regression from #902) #939 (scheduler HBM prefix-cache fix). Accuracy CI is now re-running on the rebased HEAD.

2. Performance on Qwen3 dense (27B) and other MoE (397B): cross-model benchmarks are in progress on my side. I'll post the numbers here as soon as the runs finish.

Separately, per your suggestion on #888, I've merged the fused MRoPE QK path into this PR (commit 9d0e326) and closed #888. The combined PR now delivers GDN decode + MRoPE fusion together:

• Total token throughput: 7466.90 → 8004.41 tok/s (+7.20% vs main)
• Mean E2E latency: 176401 → 164893 ms (-6.52% vs main)
• Mean TPOT: 77.44 → 71.87 ms (-7.19% vs main)
• GSM8K 5-shot: 0.895 (flex) / 0.903 (strict) — on par with main (0.8946 / 0.9052)

PR title and description have been updated to reflect the combined scope. Please take another look when you have a chance.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

[zovonoir]

@wanzhenchn gentle ping — when you have a moment, could you take another look at this PR? Quick recap of changes since your last review:

• ✅ Rebased onto main (now includes [fix](mtp): fix qwen3.5 mtp accept rate zero #966 qwen3.5 mtp fix, scheduler: fix dead HBM prefix cache under deferred-output mode (regression from #902) #939 scheduler fix, feat: support chunk prefill #740 chunk prefill, etc.)
• ✅ Merged Add fused MRoPE QK path for Qwen3.5-35B #888's fused MRoPE QK path into this PR per your suggestion (Add fused MRoPE QK path for Qwen3.5-35B #888 closed)
• ✅ Resolved the latest Copilot review item (PAD_SLOT_ID handling in the lossy decode kernel)
• ✅ Fixed an MMStar-only ConstraintViolationError → VllmBackend can only be called once regression introduced by the MRoPE path under torch.compile (early-return when compiling, eager mode unaffected)
• ✅ Lint clean (Black / Ruff / Pre Checkin all green)
• ✅ Qwen3.5-35B-A3B-FP8 TP2 accuracy CI passing on both runs
• ⏳ Cross-model benchmarks on Qwen3 dense / 397B-MoE still in progress; I'll post them as they come in

Note on the remaining red accuracy jobs (DSR1-FP8, DSV4-Pro, Qwen3.5-397B-MXFP4): these fail identically on #889 which only changes 5 lines in linear.py, so they're upstream main regressions / infra flakes, not from this PR.

Would also appreciate a look at #889 (Qwen3.5-35B-A3B PTPC support, single-file linear.py change, now rebased onto main) — it's MERGEABLE and just waiting on REVIEW_REQUIRED to clear.

Thanks!

[zovonoir]

Thanks @wanzhenchn for the offline confirmation!

CI status update for context: the latest run on 1538f12e has the target model (Qwen3.5-35B-A3B-FP8 TP2 SGLang) passing in 13m26s, along with DeepSeek-R1-0528, DeepSeek-V4-Pro / V4-Pro MTP, GLM-5.1 family, Kimi-K2.5, gpt-oss-120b, Qwen3-Next-80B, Meta-Llama-3-8B, DSR1-FP8 TP4, and Copilot review (clean). Lint (Black/Ruff/Pre Checkin) all green.

The 5 remaining red accuracy jobs are all infra (I checked each log):

Job	Real cause
Qwen3.5-35B-A3B-FP8 TP2 (vLLM)	ValueError: Free memory on device cuda:0 (176.87/251.98 GiB) ... less than 0.9 — GPU held by another process
DSR1-FP4 TP4	No such image: atom_sglang_base:ci
Kimi-K2-Thinking-MXFP4 TP4	No such image: atom_oot_base:ci
DSR1-FP8 TP8 (failed in 23s)	No such image: atom_oot_base:ci
Qwen3.5-397B-A17B-MXFP4	aiter wheel missing

These all reproduce identically on #889 (which only touches 5 lines in linear.py), so they're upstream main regressions / runner flakes, not from this PR.

@valarLip @Pleaplusone since you've both worked recently on atom/model_ops/fla_ops/ and atom/models/qwen3_next.py, would either of you have a moment to formally approve this so it can land? wanzhenchn LGTM'd offline but doesn't have approve rights.

Same ask would apply to #889 (Qwen3.5-35B-A3B PTPC support, single-file linear.py change) once this one merges.

Thanks all!

[valarLip]

the code itself is in ATOM already why we need this werid try... except