Support for Qwen3.6 (27B & 35B-A3B) Hybrid Gated DeltaNet Models on TPU v5e-8 with vLLM and SGLang

General
1

Hi vLLM Community,

I am looking to deploy the new hybrid linear-attention models from Alibaba—specifically Qwen/Qwen3.6-27B (dense) and Qwen/Qwen3.6-35B-A3B (MoE)—on a Google Cloud TPU v5e-8 instance (8 chips, 128 GB total HBM) using vLLM.

Model Architecture Overview:

These models use a 3:1 repeating hybrid sequence-mixing schedule:

  • ~75% Gated DeltaNet Layers: Linear recurrence maintaining a fixed-size recurrent state matrix updated via chunkwise delta rules.
  • ~25% Gated Softmax Attention Layers: Standard token-to-token attention requiring KV cache management (PagedAttention).
  • Target Checkpoints:
    • Qwen/Qwen3.5-27B (64 layers: 48 Gated DeltaNet + 16 Gated Attention)
    • Qwen/Qwen3.5-35B-A3B / Qwen3-Next hybrid variants

Questions for TPU / Maintainer Team:

  1. Recurrent Kernel Support in TPU Backend:
    Does the PyTorch-XLA / tpu-inference backend in vLLM currently support the chunkwise associative scan and delta-rule updates required by Gated DeltaNet (e.g., via Pallas kernels or native XLA primitives), or are these recurrent layers currently supported only via Triton on NVIDIA GPUs?

  2. Hybrid State & KV-Cache Management:
    How does vLLM handle sharding and memory allocation on TPU when mixing standard PagedAttention KV-caches with layer-wise fixed-size recurrent states across 8 TPU v5e chips?

  3. SPMD Sharding & Feasibility on v5e-8:
    Given the 128 GB HBM capacity on v5e-8 (sufficient for ~54 GB BF16 weights of the 27B model), are there any known compilation graph or tensor-parallel sharding bottlenecks when running Qwen3.5’s hybrid layout on v5e?

Any guidance, recommended flags (--tensor-parallel-size, --model-loader), or active PR/roadmap links for Qwen3.6 hybrid TPU serving would be greatly appreciated!

Thanks! cc: @Herald

2

It seems likely that vLLM’s TPU backend (PyTorch-XLA) supports basic hybrid models like Qwen3.6, but there are important caveats. The initial TPU integration in vLLM uses PyTorch-XLA and Pallas kernels for FlashAttention and PagedAttention, but support for custom recurrent kernels (like Gated DeltaNet’s chunkwise associative scan) is not explicitly documented—these may only be fully optimized on NVIDIA GPUs with Triton, and not all custom ops are available on TPU yet. See Initial TPU integration PR and TPU RFC.

For hybrid state and KV-cache management, vLLM’s architecture is designed to handle mixed attention and recurrent layers, but on TPU, sharding and memory allocation rely on PyTorch-XLA’s SPMD partitioning. The TPU backend currently supports tensor parallelism (TP), but multi-host and pipeline parallelism are still being developed. There are no explicit warnings about hybrid model sharding, but large models may hit XLA graph compilation or HBM memory limits, especially if custom recurrent ops are not fused. Recommended flags include setting --tensor-parallel-size to match the number of TPU chips and using PyTorch-XLA’s device selection. For the latest status and roadmap, see TPU RFC and TPU PR.

Would you like a step-by-step breakdown of TPU deployment for these hybrid models, or more detail on kernel support and memory layout?

Sources:

Hint: Mention @Herald in the post for followups.

3

Thanks for the detailed breakdown! Tagging @Herald for follow-up.

Given that the 75% Gated DeltaNet layers in models like Qwen/Qwen3.6-27B (dense) and Qwen/Qwen3.6-35B-A3B rely on chunkwise associative scan and delta-rule updates:

  1. Kernel Status for Qwen Hybrid Layers: Is there currently a Pallas implementation or native JAX/XLA associative scan fallback for Gated DeltaNet layers in the pipeline, or does PyTorch-XLA currently attempt unfused graph tracing when serving Qwen/Qwen3.6-27B?
  2. Compilation & Performance: If running Qwen/Qwen3.6-27B or Qwen/Qwen3.6-35B-A3B on a TPU v5e-8 instance with --tensor-parallel-size 8, has anyone tested whether these recurrent layers compile cleanly without hitting excessive XLA graph overhead, memory leaks, or compilation timeouts?
  3. Recommended Test Setup: Are there specific experimental flags, environment variables, or an active PR/branch you recommend testing to evaluate Qwen/Qwen3.6-27B and Qwen/Qwen3.6-35B-A3B on TPU v5e?

Thanks again for the pointers to the TPU RFC and PR!

4

My understanding is that, as of the latest vLLM TPU backend, there is no dedicated Pallas or native JAX/XLA kernel for Gated DeltaNet (GDN) chunkwise associative scan; PyTorch-XLA will likely trace the unfused PyTorch implementation for these layers when serving Qwen3.6-27B/35B-A3B on TPU. This may result in suboptimal performance and longer compilation times, since the recurrent scan is not fused or optimized for XLA. See Initial TPU integration PR and TPU RFC.

It seems likely that no public reports confirm clean compilation or high performance for Qwen3.6 hybrid models on TPU v5e-8 with --tensor-parallel-size 8; the RFC and PR do not mention explicit validation for these architectures. There are no documented experimental flags or PRs specifically for GDN on TPU, but using the latest vLLM main branch, PyTorch-XLA ≥2.3, and setting --tensor-parallel-size 8 is recommended. Monitoring XLA graph compilation logs and memory usage is advised. For further updates, watch the TPU RFC.

Would you like more detail on possible workarounds or how to monitor/diagnose XLA graph issues during such tests?

Sources:

Hint: Mention @Herald in the post for followups.

5

Thanks for the clear explanation! Tagging @Herald for follow-up.

Yes, I would greatly appreciate a breakdown of:

  1. Diagnostics & Logging: What environment variables (e.g., XLA_FLAGS, PT_XLA_DEBUG) or monitoring tools do you recommend to diagnose XLA graph compilation bloat and trace execution times for Qwen/Qwen3.5-27B and Qwen/Qwen3.5-35B-A3B?
  2. Workarounds: Are there any recommended runtime flags, torch compile settings, or context-length constraints (--max-model-len, --max-num-seqs) that help mitigate compilation overhead on unfused recurrent layers?
  3. Alternative Frameworks / Bridges: In the interim before Pallas kernels land in vLLM, are there recommended JAX-native pathways (or MaxText implementations) that handle Gated DeltaNet more effectively on Cloud TPU v5e?

Thanks again for the guidance!

6

It seems likely that for diagnosing XLA graph compilation and execution on TPU, you should set environment variables like XLA_FLAGS="--xla_hlo_profile --xla_dump_to=/tmp/xla_dump" and PT_XLA_DEBUG=1 to enable detailed graph dumps and profiling. Tools such as TensorBoard (with XLA traces), and inspecting /tmp/xla_dump for HLO/IR files, are recommended for analyzing graph size and execution times. See TPU RFC and PyTorch XLA docs.

For workarounds, reducing --max-model-len and --max-num-seqs can help limit graph size and compilation time. Using torch.compile with conservative dynamic shapes or disabling aggressive optimizations may also help. No vLLM-specific flags are documented for unfused recurrent layers, but keeping batch sizes and context lengths modest is generally advised.

As for alternatives, MaxText (JAX-native) is known to support Gated DeltaNet and similar architectures efficiently on TPU, as it can leverage JAX’s associative scan primitives and XLA fusion. If vLLM’s PyTorch-XLA path is too slow, trying MaxText or a JAX-based implementation may yield better performance for these hybrid models. See TPU RFC for further discussion.

Would you like a step-by-step example of setting up these diagnostics or more detail on JAX/MaxText alternatives?

Sources:

Hint: Mention @Herald in the post for followups.

7

Thanks again for the insights! @Herald

Yes, I would really appreciate both:

  1. Step-by-Step Diagnostic Setup: Could you share an example script/command demonstrating how to attach TensorBoard profiling and extract the HLO dumps on TPU v5e-8 when running Qwen/Qwen3.6-27B or Qwen/Qwen3.6-35B-A3B through vLLM?
  2. MaxText / JAX Setup & Weights: Does MaxText currently have a reference config or checkpoint conversion tool for Qwen hybrid architectures (Gated DeltaNet + Softmax Attention), or would this require custom model definition in MaxText?
  3. vLLM Roadmap: Is there an ongoing PR or tracking issue in vLLM aiming to implement Pallas-based Gated DeltaNet kernels for the TPU backend?

Thanks for all the help!

8

My understanding is that for step-by-step XLA diagnostics on TPU v5e-8, you can launch vLLM with the following environment variables to enable HLO dumps and TensorBoard profiling:

export XLA_FLAGS="--xla_hlo_profile --xla_dump_to=/tmp/xla_dump"
export PT_XLA_DEBUG=1
export XLA_TENSORBOARD_LOGDIR=/tmp/xla_tb
vllm serve Qwen/Qwen3.6-27B --tensor-parallel-size 8 ...

After running, point TensorBoard to the logdir:

tensorboard --logdir=/tmp/xla_tb

You can then inspect /tmp/xla_dump for HLO/IR files and use TensorBoard to analyze graph compilation and execution times. See PyTorch XLA debugging docs for more details.

For MaxText/JAX, it seems there is no official reference config or checkpoint converter for Qwen hybrid (Gated DeltaNet + Softmax Attention) models; using MaxText would likely require a custom model definition and manual checkpoint conversion. No direct support for Qwen hybrid architectures is documented.

Regarding vLLM’s roadmap, there is no public PR or tracking issue specifically for Pallas-based Gated DeltaNet kernels on TPU, but the TPU RFC and Initial TPU PR mention ongoing work and welcome contributions for improved kernel support.

Would you like more detail on custom model conversion for MaxText or further guidance on interpreting XLA dumps?

Sources:

Hint: Mention @Herald in the post for followups.

9

Thank you for the thorough breakdown and diagnostic commands! @Herald

This clarifies things completely:

  1. Since Gated DeltaNet lacks fused Pallas kernels in vLLM and MaxText lacks out-of-the-box model definitions, running on TPU v5e via PyTorch-XLA will require profiling unfused execution.
  2. I will run the benchmark on our v5e-8 instance using the provided XLA_FLAGS and XLA_TENSORBOARD_LOGDIR setup to inspect compilation times and memory bottlenecks in the HLO dump.

Could you share 1–2 key metrics or HLO pattern markers we should look for in /tmp/xla_dump to identify if recurrent scan loops are causing excessive recompilation or graph bloating?

Thanks again for all the help!

10

It could be possible that the most important HLO patterns to watch for are: (1) repeated or very large while/scan/fold operations in the HLO graph, which indicate unfused recurrent scan loops, and (2) excessive graph size or many unique HLO modules, which can signal recompilation or graph bloat. You might also check for high memory allocation in the HLO profile, especially for buffers associated with recurrent state updates. See PyTorch XLA debugging docs for more on interpreting HLO dumps.

Would you like more detail on how to interpret specific HLO op names or memory metrics in the dump?

Sources:

Hint: Mention @Herald in the post for followups.

11

Thank you so much for all the detailed guidance throughout this thread! @Herald

This gives us a clear testing plan:

  1. We will benchmark Qwen/Qwen3.6-27B on the TPU v5e-8 instance with --tensor-parallel-size 8.
  2. We’ll monitor /tmp/xla_dump and TensorBoard specifically for nested while/recurrent loops and buffer allocation churn.

I really appreciate the time and pointers to the docs!

12

You’re very welcome! Monitoring for nested while/scan ops and buffer churn in /tmp/xla_dump and TensorBoard is the right approach for diagnosing recurrent layer bottlenecks on TPU. If you encounter specific HLO patterns or memory issues, sharing those logs can help the community provide targeted advice.

If you need further help interpreting XLA dumps or run into TPU-specific issues with Qwen3.6-27B, feel free to follow up. Good luck with your benchmarking!

Sources:

Hint: Mention @Herald in the post for followups.