Use share_mutable_buffers to eliminate select_scatter overhead

Revert from explicit state passing back to registered buffers with
in-place updates (KVCache, conv_state, recurrent_state). Export with
share_mutable_buffers=True so both prefill and forward methods share
mutable state via mem_id=2. C++ runner uses share_memory_arenas=true
and only passes (tokens, input_pos) — no CUDA runtime dependency.

Results: 84.5 tok/s (up from 77.4), 0 select_scatter ops in profile,
65 D2H memcpy (logits only).

Author: gasoonjia

examples/models/qwen3_5_moe/CMakeLists.txt

@@ -44,7 +44,7 @@ list(
 
 # CUDA backend (required)
 find_package(CUDAToolkit REQUIRED)
-list(APPEND link_libraries aoti_cuda_backend CUDA::cudart)
+list(APPEND link_libraries aoti_cuda_backend)
 executorch_target_link_options_shared_lib(aoti_cuda_backend)
 
 # Tokenizer

examples/models/qwen3_5_moe/export.py

@@ -109,6 +109,9 @@ def load_prequantized_model(prequantized_dir, max_seq_len=4096):
     runtime_prefixes = (
         ".mask",
         ".inv_freq",
+        ".kv_cache.",
+        ".conv_state",
+        ".recurrent_state",
     )
     expected_missing = {k for k in missing if any(p in k for p in runtime_prefixes)}
     weight_missing = set(missing) - expected_missing
@@ -310,7 +313,8 @@ def _materialize_buffers(model, config):
 
     Replaces meta buffers with real tensors on CPU, recomputes RoPE
     inv_freq and causal masks. State buffers (KV cache, conv/recurrent
-    state) are no longer registered buffers — they are explicit function args.
+    state) are zero-initialized registered buffers that will be shared
+    across methods via share_mutable_buffers.
     """
     # Masks stay bool, inv_freq stays float32.
     for fqn, buf in list(model.named_buffers()):
@@ -359,8 +363,9 @@ def export_and_lower(model, config, args):
       - "prefill": prefill path (T>=2), uses chunked FLA triton_op with
         dynamic sequence length.
 
-    Both methods take explicit state tensors (conv_states, recurrent_states,
-    k_caches, v_caches) as inputs and return updated state as outputs.
+    Both methods share mutable state buffers (KV cache, conv_state,
+    recurrent_state) via share_mutable_buffers=True. The model uses
+    registered buffers with in-place updates — no state in/out args.
     """
     import torch._inductor.config as inductor_config
 
@@ -381,19 +386,14 @@ def export_and_lower(model, config, args):
     # -O0 compiles ~8x faster than -O1 with no measurable runtime impact.
     inductor_config.aot_inductor.compile_wrapper_opt_level = "O0"
 
-    # Create initial state tensors
-    conv_states, recurrent_states, k_caches, v_caches = \
-        Qwen35MoE.make_initial_state(config)
-
     # --- Decode method (T=1, static shape) ---
     print("Exporting decode method (forward)...")
     decode_tokens = torch.tensor([[0]], dtype=torch.long)
     decode_pos = torch.tensor([0], dtype=torch.long)
     with torch.no_grad():
         decode_ep = export(
             model,
-            (decode_tokens, decode_pos,
-             conv_states, recurrent_states, k_caches, v_caches),
+            (decode_tokens, decode_pos),
             strict=True,
         )
     print("Decode export successful!")
@@ -403,21 +403,14 @@ def export_and_lower(model, config, args):
     prefill_tokens = torch.tensor([[0, 1]], dtype=torch.long)
     prefill_pos = torch.tensor([0, 1], dtype=torch.long)
     seq_dim = Dim("seq_len", min=2, max=config.max_seq_len - 1)
-    # Dynamic shapes: only tokens dim 1 and pos dim 0 are dynamic;
-    # state tensors have static shapes.
     prefill_dynamic_shapes = (
         {1: seq_dim},   # tokens
         {0: seq_dim},   # input_pos
-        None,           # conv_states
-        None,           # recurrent_states
-        None,           # k_caches
-        None,           # v_caches
     )
     with torch.no_grad():
         prefill_ep = export(
             model,
-            (prefill_tokens, prefill_pos,
-             conv_states, recurrent_states, k_caches, v_caches),
+            (prefill_tokens, prefill_pos),
             dynamic_shapes=prefill_dynamic_shapes,
             strict=True,
         )
@@ -426,30 +419,13 @@ def export_and_lower(model, config, args):
     # Lower with CUDA backend (per-method partitioners to avoid so_blob collision)
     print("Lowering to ExecuTorch with CUDA...")
 
-    num_fla = sum(1 for t in config.layer_types if t == "linear_attention")
-    num_attn = sum(1 for t in config.layer_types if t == "full_attention")
-    conv_dim = (
-        config.linear_num_key_heads * config.linear_key_head_dim * 2
-        + config.linear_num_value_heads * config.linear_value_head_dim
-    )
-
     metadata = {
         "get_max_seq_len": config.max_seq_len,
         "get_vocab_size": config.vocab_size,
         "get_n_layers": config.num_hidden_layers,
         "use_kv_cache": True,
         "use_sdpa_with_kv_cache": False,
         "enable_dynamic_shape": True,
-        # State shape metadata for C++ runner
-        "get_num_fla_layers": num_fla,
-        "get_num_attn_layers": num_attn,
-        "get_conv_dim": conv_dim,
-        "get_conv_kernel_size": config.linear_conv_kernel_dim,
-        "get_num_v_heads": config.linear_num_value_heads,
-        "get_head_k_dim": config.linear_key_head_dim,
-        "get_head_v_dim": config.linear_value_head_dim,
-        "get_n_kv_heads": config.num_kv_heads,
-        "get_head_dim": config.head_dim,
     }
     et_prog = to_edge_transform_and_lower(
         {"forward": decode_ep, "prefill": prefill_ep},
@@ -471,7 +447,11 @@ def export_and_lower(model, config, args):
         config=ExecutorchBackendConfig(
             extract_delegate_segments=True,
             do_quant_fusion_and_const_prop=True,
-            memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
+            memory_planning_pass=MemoryPlanningPass(
+                alloc_graph_input=False,
+                share_mutable_buffers=True,
+            ),
+            emit_mutable_buffer_names=True,
         ),
     )
 

examples/models/qwen3_5_moe/inference.py

@@ -71,34 +71,24 @@ def _move_to_cuda(model, config):
 
 
 def generate(
-    model, config, tokenizer, prompt, max_new_tokens=128, temperature=0.0, eos_token_ids=None
+    model, tokenizer, prompt, max_new_tokens=128, temperature=0.0, eos_token_ids=None
 ):
-    """Generate text autoregressively with explicit state passing.
+    """Generate text autoregressively.
 
-    Prefills one token at a time (the chunk_gated_delta_rule kernel's chunked
-    path has numerical issues with T>1 in eager mode; token-by-token uses the
-    stable recurrent path).
+    State (KV cache, conv_state, recurrent_state) is managed internally
+    via registered buffers — the model signature is just (tokens, input_pos).
     """
-    from executorch.examples.models.qwen3_5_moe.model import Qwen35MoE
-
     if eos_token_ids is None:
         eos_token_ids = set()
 
     input_ids = tokenizer.encode(prompt).ids
 
-    # Initialize state tensors
-    conv_states, recurrent_states, k_caches, v_caches = (
-        Qwen35MoE.make_initial_state(config, dtype=torch.bfloat16, device="cuda")
-    )
-
-    # Prefill: one token at a time
+    # Prefill: one token at a time (recurrent path is stable for T=1)
     with torch.no_grad():
         for i, tok_id in enumerate(input_ids):
             tok = torch.tensor([[tok_id]], dtype=torch.long, device="cuda")
             pos = torch.tensor([i], dtype=torch.long, device="cuda")
-            logits, conv_states, recurrent_states, k_caches, v_caches = model(
-                tok, pos, conv_states, recurrent_states, k_caches, v_caches
-            )
+            logits = model(tok, pos)
 
     # Sample first generated token
     next_token = _sample(logits[:, -1, :], temperature)
@@ -109,10 +99,7 @@ def generate(
     with torch.no_grad():
         for i in range(max_new_tokens - 1):
             pos = torch.tensor([seq_len + i], device="cuda")
-            logits, conv_states, recurrent_states, k_caches, v_caches = model(
-                next_token.unsqueeze(0), pos,
-                conv_states, recurrent_states, k_caches, v_caches
-            )
+            logits = model(next_token.unsqueeze(0), pos)
             next_token = _sample(logits[:, -1, :], temperature)
             tok_id = next_token.item()
             generated.append(tok_id)
@@ -193,7 +180,6 @@ def main():
     t0 = time.perf_counter()
     output = generate(
         model,
-        config,
         tokenizer,
         args.prompt,
         max_new_tokens=args.max_new_tokens,