feat(RL): add RL support for verl#1298
Open
shihaobai wants to merge 321 commits into
Open
Conversation
Implement comprehensive support for Qwen3Next model with linear attention mechanism: Model Features: - Implement linear attention with MTP (Multi-Token Prediction) capability - Add custom Triton kernels for gated delta networks (GDN) operations - Support chunked operations for efficient attention computation - Add specialized buffer pool and memory managers for linear attention Triton Kernels: - Add causal_conv1d for efficient convolution operations - Implement chunk-based operations (chunk_o, chunk_delta_h, chunk_scaled_dot_kkt) - Add gated delta network kernels (fused_gdn_gating, gdn_decode_mtp) - Implement fused normalization (gemma_rmsnorm, gated_rmsnorm) Infrastructure: - Add hybrid radix cache for efficient memory management - Implement mamba cache manager for state management - Add allocator utilities for buffer management - Add parameter weight abstraction for flexible weight handling - Update model registration and API endpoints Performance Optimizations: - Add H200 autotune configurations for all Triton kernels - Optimize memory allocation with custom kernels - Support chunked prefill and decode backends This implementation enables efficient inference for models with linear attention mechanisms, providing significant speedup for long sequence lengths.
- Reduce Triton kernels from 6 (1D/2D/3D × p2p/broadcast) to 2 (1D only) by flattening contiguous trailing dimensions via tensor view - Wire up MambaCacheManager to use the Triton kernels instead of PyTorch advanced indexing with Python for-loops - Cast strides to int64 in kernels to prevent pointer arithmetic overflow - Add Qwen3.5 multimodal vision-language model support
- Add mamba_cache_ratio parameter (default 0.5) - Change mamba_cache_size default from 3000 to None - Implement automatic memory allocation based on ratio - Add clear error messages with solutions when memory insufficient - Maintain backward compatibility with explicit mamba_cache_size Ratio formula: mamba_memory = total_available * ratio / (1 + ratio) - ratio=0.5 -> 33% mamba, 67% KV - ratio=1.0 -> 50% mamba, 50% KV - ratio=2.0 -> 67% mamba, 33% KV
Change ratio meaning from complex formula to simple percentage: - Old: ratio = mamba / kv, mamba = total * ratio / (1+ratio) - New: ratio = mamba / total, mamba = total * ratio This makes the ratio more intuitive: - 0.3 → 30% mamba, 70% KV - 0.5 → 50% mamba, 50% KV (default) - 0.7 → 70% mamba, 30% KV Also simplifies error message recommendation formula.
- Replace get_radix_cache_class() classmethod with radix_cache_class class attribute in TpPartBaseModel and Qwen3NextTpPartModel - Move RadixCache/HybridRadixCache imports to module top-level - Update base_backend.py to access radix_cache_class directly - Replace alloc_buffer_for_req_triton with simpler indexed PyTorch assignment - Remove now-unused alloc_buffer_kernel.py Triton kernel - Revert LOADWORKER default to 1 and remove language_model. prefix stripping
The sgl_kernel.fwd.default API requires attention_chunk before softcap. This file was missed when the parameter was added in commit a4ab210. Also update sgl-kernel from 0.3.7.post1 to 0.3.21 which supports this API.
- Rename copy_buffer_p2p → copy_mamba_buffer (indexed 1:1 slot copy) - Rename copy_buffer_broadcast → fork_mamba_buffer (1:N MTP fork) - Unify chunk offset param name (pair_idx_offset/copy_idx_offset → chunk_offset) - Rename stride_index → stride_slot to reflect the slot/cache dimension - Rename src_idx_in_batch → src_chunk_idx in fork kernel - Extract _MAX_GRID_DIM = 65535 module constant (was duplicated inline) - Add divisibility assertion before implicit // in fork autotuned wrapper - Update autotuner cache keys to match new names
Resolved conflicts by preserving RL-specific features from HEAD (routing capture, Neo models, weight versioning, RL args) while integrating Qwen3.5/Qwen3Next/Qwen3Omni models and None-safe sampling param fixes from qwen3.5_clean. Conflict resolution decisions: - models/__init__.py: keep both Neo and Qwen3.5/Omni imports - qwen3_moe/model.py: add conditional num_experts guard + keep routing capture - api_start.py: keep HEAD's _launch_subprocesses refactor structure - sampling_params.py: keep HEAD's _get/_cfg helpers (cleaner than per-field None checks) - start_args_type.py: merge all fields from both branches - infer_batch.py: keep routing data extraction + use _free_req_mem_and_buffers - base_backend.py: keep HEAD's utils imports + add communication_op imports - visualserver/model_rpc.py: keep both Neo and Qwen3Omni branches
KVCACHE_TOKEN_CAN_USE_NUM_SHM_NAME and MAMBA_CACHE_CAN_USE_NUM_SHM_NAME were evaluated at module import time, calling get_unique_server_name() before set_unique_server_name(args) had been called. This caused: TypeError: argument of type 'NoneType' is not iterable Replace module-level constants with _get_*_shm_name() functions so the env var lookup is deferred until the values are actually needed (inside __init__ / class methods, after server startup).
Two critical bugs introduced by merge origin/qwen3.5_clean: 1. mm_slicer.py: Weight slicing dimension changed from shape[0]/shape[1] to shape[-2]/shape[-1], which corrupts 2D tensor weights during TP. For 2D weights [out_dim, in_dim], slicing dim -2 instead of dim 0 produces incorrect results with tensor parallelism. 2. rmsnorm.py: Precision regression in Triton kernel: - Weight loaded without float32 conversion - x_hat computed then cast to bfloat16 before multiplication - This loses precision during normalization, degrading accuracy Both issues caused intermittent garbage output and accuracy drop (0.88 -> 0.65) for qwen3.5_moe model.
Implement comprehensive support for Qwen3Next model with linear attention mechanism: Model Features: - Implement linear attention with MTP (Multi-Token Prediction) capability - Add custom Triton kernels for gated delta networks (GDN) operations - Support chunked operations for efficient attention computation - Add specialized buffer pool and memory managers for linear attention Triton Kernels: - Add causal_conv1d for efficient convolution operations - Implement chunk-based operations (chunk_o, chunk_delta_h, chunk_scaled_dot_kkt) - Add gated delta network kernels (fused_gdn_gating, gdn_decode_mtp) - Implement fused normalization (gemma_rmsnorm, gated_rmsnorm) Infrastructure: - Add hybrid radix cache for efficient memory management - Implement mamba cache manager for state management - Add allocator utilities for buffer management - Add parameter weight abstraction for flexible weight handling - Update model registration and API endpoints Performance Optimizations: - Add H200 autotune configurations for all Triton kernels - Optimize memory allocation with custom kernels - Support chunked prefill and decode backends This implementation enables efficient inference for models with linear attention mechanisms, providing significant speedup for long sequence lengths.
- Reduce Triton kernels from 6 (1D/2D/3D × p2p/broadcast) to 2 (1D only) by flattening contiguous trailing dimensions via tensor view - Wire up MambaCacheManager to use the Triton kernels instead of PyTorch advanced indexing with Python for-loops - Cast strides to int64 in kernels to prevent pointer arithmetic overflow - Add Qwen3.5 multimodal vision-language model support
- Add mamba_cache_ratio parameter (default 0.5) - Change mamba_cache_size default from 3000 to None - Implement automatic memory allocation based on ratio - Add clear error messages with solutions when memory insufficient - Maintain backward compatibility with explicit mamba_cache_size Ratio formula: mamba_memory = total_available * ratio / (1 + ratio) - ratio=0.5 -> 33% mamba, 67% KV - ratio=1.0 -> 50% mamba, 50% KV - ratio=2.0 -> 67% mamba, 33% KV
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.