Vllm modelling cleanup

[oleksost]

Context:

• 'oo/feature/vllm-aprie2-model-permixergraph/` branch was used for benchmarking throughout and obtaining some of the results for super Apriel (production tasks, i.e. non dev)
• this is cleanup of the modeling file, which also includes some of the throughout related fixes (mostly relevant for cost model) described below

Tests (vllm 0.14rc1):

• serving as full supernet using layer palcement code works well
• serving in "surgery" mode, i.e. only predefined layout with plain vllm serve works well
• currently making sure that we can reproduce the throughout numbers

Summary

Three categories of bug fixes, a KV cache throughput fix, and a general cleanup of the modeling file.

Bug fixes

• GDN/KDA OOB crash during CUDA graph capture (9a6562d): During FULL CUDA graph capture vLLM creates dummy runs with batch sizes up to
  max_cudagraph_capture_size. When gpu_memory_utilization is low, the number of KV cache blocks can be smaller than the capture batch size, causing
  causal_conv1d_update and recurrent kernel assertions to fail. Fix: clamp num_actual_tokens to conv_state.size(0) only during CUDA graph capture
  (cudagraph_capturing_enabled). Also fixes cu_seqlens and ssm_state_indices slicing for decode paths in both GDN and KDA, and removes a spurious
  return-value assignment for the in-place causal_conv1d_update.
• cudagraph_capturing_enabled import-by-value bug causing OOB crash on re-prefill (60541b0): from vllm.compilation.monitor import
  cudagraph_capturing_enabled captures the default value True at import time. When vLLM later calls set_cudagraph_capturing_enabled(False), the
  module-level name stays True, causing num_actual_tokens to be wrongly clamped during normal prefill — truncating inputs while query_start_loc still
  referenced the full batch, crashing causal_conv1d_fn. Fix: use vllm.compilation.monitor.cudagraph_capturing_enabled (module attribute access)
  everywhere.

KV cache grouping fix

• Throughput regression fix for hybrid models with singleton mixer types (7a948a4): vLLM's _get_kv_cache_groups_uniform_page_size uses
  min(num_layers_per_type) as group size, creating O(num_layers) groups for models like a12_g1_k11 (where min=1). Each group triggers a metadata rebuild
  per forward step, causing >2× throughput regression. Monkey-patches the grouping function to use max(num_layers_per_type) when a
  singleton/near-singleton type is detected, reducing groups to O(num_types).
• apriel2_gdn_attention_core always registered as PIECEWISE splitting op: Without the graph break, the decode path gets baked into a compiled prefill
  piece — prefill batches then replay the wrong path causing illegal memory access.
• Apriel2Attention.get_kv_cache_spec: Removes spurious else: branch (cosmetic, same behavior).
• fused_gdn_gating_kernel: Removes tl.constexpr from total_elements (was causing issues with variable batch sizes).

Cleanup

• Removed ~800 lines of WIP code: per-mixer CUDA graph caching (MixerGraphCache), CPU offload for inactive mixers, APRIEL2_FULL_CUDA_GRAPHS /
  APRIEL2_MIXER_CUDA_GRAPHS / APRIEL2_OFFLOAD_INACTIVE_MIXERS env vars, and associated documentation.
• Removed all debug artifacts: DEBUG_* flags, _debug_print/_debug_tensor/_debug_state_stats methods, commented-out debug calls, and CUDA sync debug
  points.