diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py index 0329d110244c..a8b1c54780bd 100644 --- a/benchmarks/attention_benchmarks/benchmark.py +++ b/benchmarks/attention_benchmarks/benchmark.py @@ -47,6 +47,8 @@ is_mla_backend, ) +from vllm.v1.worker.workspace import init_workspace_manager + def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: """Run standard attention benchmark (Flash/Triton/FlashInfer).""" @@ -462,7 +464,7 @@ def main(): parser.add_argument( "--batch-specs", nargs="+", - default=["q2k", "8q1s1k"], + default=None, help="Batch specifications using extended grammar", ) @@ -478,6 +480,21 @@ def main(): parser.add_argument("--repeats", type=int, default=1, help="Repetitions") parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations") parser.add_argument("--profile-memory", action="store_true", help="Profile memory") + parser.add_argument( + "--kv-cache-dtype", + default="auto", + choices=["auto", "fp8"], + help="KV cache dtype: auto or fp8", + ) + parser.add_argument( + "--cuda-graphs", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Launch kernels with CUDA graphs to eliminate CPU overhead" + "in measurements (default: True)" + ), + ) # Parameter sweep (use YAML config for advanced sweeps) parser.add_argument( @@ -536,21 +553,24 @@ def main(): # Batch specs and sizes # Support both explicit batch_specs and generated batch_spec_ranges - if "batch_spec_ranges" in yaml_config: - # Generate batch specs from ranges - generated_specs = generate_batch_specs_from_ranges( - yaml_config["batch_spec_ranges"] - ) - # Combine with any explicit batch_specs - if "batch_specs" in yaml_config: - args.batch_specs = yaml_config["batch_specs"] + generated_specs - else: - args.batch_specs = generated_specs - console.print( - f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]" - ) - elif "batch_specs" in yaml_config: - args.batch_specs = yaml_config["batch_specs"] + # CLI --batch-specs takes precedence over YAML when provided. + cli_batch_specs_provided = args.batch_specs is not None + if not cli_batch_specs_provided: + if "batch_spec_ranges" in yaml_config: + # Generate batch specs from ranges + generated_specs = generate_batch_specs_from_ranges( + yaml_config["batch_spec_ranges"] + ) + # Combine with any explicit batch_specs + if "batch_specs" in yaml_config: + args.batch_specs = yaml_config["batch_specs"] + generated_specs + else: + args.batch_specs = generated_specs + console.print( + f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]" + ) + elif "batch_specs" in yaml_config: + args.batch_specs = yaml_config["batch_specs"] if "batch_sizes" in yaml_config: args.batch_sizes = yaml_config["batch_sizes"] @@ -575,6 +595,10 @@ def main(): args.warmup_iters = yaml_config["warmup_iters"] if "profile_memory" in yaml_config: args.profile_memory = yaml_config["profile_memory"] + if "kv_cache_dtype" in yaml_config: + args.kv_cache_dtype = yaml_config["kv_cache_dtype"] + if "cuda_graphs" in yaml_config: + args.cuda_graphs = yaml_config["cuda_graphs"] # Parameter sweep configuration if "parameter_sweep" in yaml_config: @@ -629,12 +653,18 @@ def main(): # Determine backends backends = args.backends or ([args.backend] if args.backend else ["flash"]) prefill_backends = getattr(args, "prefill_backends", None) + if not args.batch_specs: + args.batch_specs = ["q2k", "8q1s1k"] console.print(f"Backends: {', '.join(backends)}") if prefill_backends: console.print(f"Prefill backends: {', '.join(prefill_backends)}") console.print(f"Batch specs: {', '.join(args.batch_specs)}") + console.print(f"KV cache dtype: {args.kv_cache_dtype}") + console.print(f"CUDA graphs: {args.cuda_graphs}") console.print() + init_workspace_manager(args.device) + # Run benchmarks all_results = [] @@ -687,6 +717,8 @@ def main(): repeats=args.repeats, warmup_iters=args.warmup_iters, profile_memory=args.profile_memory, + kv_cache_dtype=args.kv_cache_dtype, + use_cuda_graphs=args.cuda_graphs, ) # Add decode pipeline config @@ -839,6 +871,8 @@ def main(): "repeats": args.repeats, "warmup_iters": args.warmup_iters, "profile_memory": args.profile_memory, + "kv_cache_dtype": args.kv_cache_dtype, + "use_cuda_graphs": args.cuda_graphs, } all_results = run_model_parameter_sweep( backends, @@ -861,6 +895,8 @@ def main(): "repeats": args.repeats, "warmup_iters": args.warmup_iters, "profile_memory": args.profile_memory, + "kv_cache_dtype": args.kv_cache_dtype, + "use_cuda_graphs": args.cuda_graphs, } all_results = run_parameter_sweep( backends, args.batch_specs, base_config_args, args.parameter_sweep, console @@ -891,6 +927,8 @@ def main(): repeats=args.repeats, warmup_iters=args.warmup_iters, profile_memory=args.profile_memory, + kv_cache_dtype=args.kv_cache_dtype, + use_cuda_graphs=args.cuda_graphs, ) result = run_benchmark(config) diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py index 208d6273c928..74d9e239725d 100644 --- a/benchmarks/attention_benchmarks/common.py +++ b/benchmarks/attention_benchmarks/common.py @@ -213,6 +213,9 @@ class BenchmarkConfig: profile_memory: bool = False use_cuda_graphs: bool = False + # "auto" or "fp8" + kv_cache_dtype: str = "auto" + # MLA-specific prefill_backend: str | None = None kv_lora_rank: int | None = None @@ -369,6 +372,7 @@ def save_csv(self, results: list[BenchmarkResult], path: str): "backend", "batch_spec", "num_layers", + "kv_cache_dtype", "mean_time", "std_time", "throughput", @@ -382,6 +386,7 @@ def save_csv(self, results: list[BenchmarkResult], path: str): "backend": r.config.backend, "batch_spec": r.config.batch_spec, "num_layers": r.config.num_layers, + "kv_cache_dtype": r.config.kv_cache_dtype, "mean_time": r.mean_time, "std_time": r.std_time, "throughput": r.throughput_tokens_per_sec or 0, diff --git a/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml b/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml index b555d90cbf62..c342e9fb8c1a 100644 --- a/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml +++ b/benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml @@ -30,9 +30,9 @@ batch_specs: - "2q16k_32q1s4k" # 2 very large prefill + 32 decode # Context extension + decode - - "2q1kkv2k_16q1s1k" # 2 extend + 16 decode - - "4q2kkv4k_32q1s2k" # 4 extend + 32 decode - - "2q1kkv8k_32q1s2k" # 2 large extend + 32 decode + - "2q1ks2k_16q1s1k" # 2 extend + 16 decode + - "4q2ks4k_32q1s2k" # 4 extend + 32 decode + - "2q1ks8k_32q1s2k" # 2 large extend + 32 decode # Explicitly chunked prefill - "q8k" # 8k prefill with chunking hint diff --git a/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml b/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml new file mode 100644 index 000000000000..689c9f3c3c66 --- /dev/null +++ b/benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml @@ -0,0 +1,58 @@ +# MLA decode-only benchmark configuration + +model: + name: "deepseek-v3" + num_layers: 60 + num_q_heads: 128 # Base value, can be swept for TP simulation + num_kv_heads: 1 # MLA uses single latent KV + head_dim: 576 + kv_lora_rank: 512 + qk_nope_head_dim: 128 + qk_rope_head_dim: 64 + v_head_dim: 128 + block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128 + +# Model parameter sweep: simulate tensor parallelism by varying num_q_heads +# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads +model_parameter_sweep: + param_name: "num_q_heads" + values: [128, 64, 32, 16] + label_format: "{backend}_{value}h" + +batch_specs: + # Small batches, varying sequence lengths + - "16q1s512" # 16 requests, 512 KV cache + - "16q1s1k" # 16 requests, 1k KV cache + - "16q1s2k" # 16 requests, 2k KV cache + - "16q1s4k" # 16 requests, 4k KV cache + + # Medium batches + - "32q1s1k" # 32 requests, 1k KV cache + - "32q1s2k" # 32 requests, 2k KV cache + - "32q1s4k" # 32 requests, 4k KV cache + - "32q1s8k" # 32 requests, 8k KV cache + + # Large batches + - "64q1s1k" # 64 requests, 1k KV cache + - "64q1s2k" # 64 requests, 2k KV cache + - "64q1s4k" # 64 requests, 4k KV cache + - "64q1s8k" # 64 requests, 8k KV cache + + # Very large batches + - "128q1s1k" # 128 requests, 1k KV cache + - "128q1s2k" # 128 requests, 2k KV cache + - "128q1s4k" # 128 requests, 4k KV cache + - "128q1s8k" # 128 requests, 8k KV cache + + # Long context + - "32q1s16k" # 32 requests, 16k KV cache + - "32q1s32k" # 32 requests, 32k KV cache + +backends: + - FLASHMLA_SPARSE + - FLASHINFER_MLA_SPARSE + +device: "cuda:0" +repeats: 100 +warmup_iters: 10 +profile_memory: true diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py index 0d612e374a12..f8bc7b4a10ed 100644 --- a/benchmarks/attention_benchmarks/mla_runner.py +++ b/benchmarks/attention_benchmarks/mla_runner.py @@ -60,9 +60,11 @@ def create_minimal_vllm_config( model_name: str = "deepseek-v3", block_size: int = 128, max_num_seqs: int = 256, + max_num_batched_tokens: int = 8192, mla_dims: dict | None = None, index_topk: int | None = None, prefill_backend: str | None = None, + kv_cache_dtype: str = "auto", ) -> VllmConfig: """ Create minimal VllmConfig for MLA benchmarks. @@ -149,13 +151,13 @@ def create_minimal_vllm_config( cache_config = CacheConfig( block_size=block_size, gpu_memory_utilization=0.9, - cache_dtype="auto", + cache_dtype=kv_cache_dtype, enable_prefix_caching=False, ) scheduler_config = SchedulerConfig( max_num_seqs=max_num_seqs, - max_num_batched_tokens=8192, + max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs), max_model_len=32768, is_encoder_decoder=False, enable_chunked_prefill=True, @@ -535,6 +537,7 @@ def _create_backend_impl( device: torch.device, max_num_tokens: int = 8192, index_topk: int | None = None, + kv_cache_dtype: str = "auto", ): """ Create backend implementation instance. @@ -583,7 +586,7 @@ def _create_backend_impl( "num_kv_heads": mla_dims["num_kv_heads"], "alibi_slopes": None, "sliding_window": None, - "kv_cache_dtype": "auto", + "kv_cache_dtype": kv_cache_dtype, "logits_soft_cap": None, "attn_type": "decoder", "kv_sharing_target_layer_name": None, @@ -701,6 +704,7 @@ def _run_single_benchmark( mla_dims: dict, device: torch.device, indexer=None, + kv_cache_dtype: str | None = None, ) -> BenchmarkResult: """ Run a single benchmark iteration. @@ -734,49 +738,124 @@ def _run_single_benchmark( ) # Create KV cache - kv_cache = torch.zeros( - num_blocks, - block_size, - mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"], - device=device, - dtype=torch.bfloat16, - ) + if kv_cache_dtype is None: + kv_cache_dtype = getattr(config, "kv_cache_dtype", "auto") + head_size = mla_dims["kv_lora_rank"] + mla_dims["qk_rope_head_dim"] + if kv_cache_dtype == "fp8_ds_mla": + # FlashMLA sparse custom format: 656 bytes per token, stored as uint8. + # Layout: kv_lora_rank fp8 bytes + 4 float32 tile scales + # + 2*rope_dim bf16 bytes + # = 512 + 16 + 128 = 656 bytes for DeepSeek dims. + kv_cache = torch.zeros( + num_blocks, + block_size, + 656, + device=device, + dtype=torch.uint8, + ) + elif kv_cache_dtype == "fp8": + from vllm.platforms import current_platform - # Create input tensors for both decode and prefill modes - decode_inputs, prefill_inputs = _create_input_tensors( - total_q, - mla_dims, - backend_cfg["query_format"], - device, - torch.bfloat16, - ) + kv_cache = torch.zeros( + num_blocks, + block_size, + head_size, + device=device, + dtype=torch.uint8, + ).view(current_platform.fp8_dtype()) + else: + kv_cache = torch.zeros( + num_blocks, + block_size, + head_size, + device=device, + dtype=torch.bfloat16, + ) # Fill indexer with random indices for sparse backends is_sparse = backend_cfg.get("is_sparse", False) if is_sparse and indexer is not None: indexer.fill_random_indices(total_q, max_kv_len) - # Determine which forward method to use based on metadata - if metadata.decode is not None: - forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer) - elif metadata.prefill is not None: - forward_fn = lambda: impl.forward_mha( - prefill_inputs["q"], - prefill_inputs["k_c_normed"], - prefill_inputs["k_pe"], - kv_cache, - metadata, - prefill_inputs["k_scale"], - prefill_inputs["output"], - ) - else: + # Determine which forward methods to use based on metadata. + # Sparse MLA backends always use forward_mqa + has_decode = is_sparse or getattr(metadata, "decode", None) is not None + has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None + if not has_decode and not has_prefill: raise RuntimeError("Metadata has neither decode nor prefill metadata") + num_decode = ( + metadata.num_decode_tokens + if (has_decode and has_prefill) + else total_q + if has_decode + else 0 + ) + num_prefill = total_q - num_decode + + # Some backends requires fp8 queries when using fp8 KV cache. + is_fp8_kvcache = kv_cache_dtype.startswith("fp8") + quantize_query = is_fp8_kvcache and getattr( + impl, "supports_quant_query_input", False + ) + + # quantize_query forces concat format + query_fmt = "concat" if quantize_query else backend_cfg["query_format"] + + # Create decode query tensors + if has_decode: + decode_inputs, _ = _create_input_tensors( + num_decode, mla_dims, query_fmt, device, torch.bfloat16 + ) + # Cast decode query to fp8 if the backend supports it + if quantize_query: + from vllm.platforms import current_platform + + if isinstance(decode_inputs, tuple): + decode_inputs = torch.cat(list(decode_inputs), dim=-1) + decode_inputs = decode_inputs.to(current_platform.fp8_dtype()) + + # Create prefill input tensors + if has_prefill: + _, prefill_inputs = _create_input_tensors( + num_prefill, mla_dims, query_fmt, device, torch.bfloat16 + ) + + # Build forward function + def forward_fn(): + results = [] + if has_decode: + results.append(impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)) + if has_prefill: + results.append( + impl.forward_mha( + prefill_inputs["q"], + prefill_inputs["k_c_normed"], + prefill_inputs["k_pe"], + kv_cache, + metadata, + prefill_inputs["k_scale"], + prefill_inputs["output"], + ) + ) + return results[0] if len(results) == 1 else tuple(results) + # Warmup for _ in range(config.warmup_iters): forward_fn() torch.accelerator.synchronize() + # Optionally capture a CUDA graph after warmup. + # Graph replay eliminates CPU launch overhead so timings reflect pure + # kernel time. + if config.use_cuda_graphs: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward_fn() + benchmark_fn = graph.replay + else: + benchmark_fn = forward_fn + # Benchmark times = [] for _ in range(config.repeats): @@ -785,7 +864,7 @@ def _run_single_benchmark( start.record() for _ in range(config.num_layers): - forward_fn() + benchmark_fn() end.record() torch.accelerator.synchronize() @@ -852,13 +931,30 @@ def _run_mla_benchmark_batched( # Determine if this is a sparse backend is_sparse = backend_cfg.get("is_sparse", False) + # Extract kv_cache_dtype from the first config + kv_cache_dtype = getattr(first_config, "kv_cache_dtype", "auto") + + # FlashMLA sparse only supports "fp8_ds_mla" internally (not generic "fp8"). + # Remap here so the user can pass --kv-cache-dtype fp8 regardless of backend. + if backend.upper() == "FLASHMLA_SPARSE" and kv_cache_dtype == "fp8": + kv_cache_dtype = "fp8_ds_mla" + + # Compute max total_q across all configs so the metadata builder buffer + # and scheduler config are large enough for all batch specs. + max_total_q = max( + sum(r.q_len for r in parse_batch_spec(cfg.batch_spec)) + for cfg, *_ in configs_with_params + ) + # Create and set vLLM config for MLA (reused across all benchmarks) vllm_config = create_minimal_vllm_config( model_name="deepseek-v3", # Used only for model path block_size=block_size, + max_num_batched_tokens=max_total_q, mla_dims=mla_dims, # Use custom dims from config or default index_topk=index_topk if is_sparse else None, prefill_backend=prefill_backend, + kv_cache_dtype=kv_cache_dtype, ) results = [] @@ -883,7 +979,9 @@ def _run_mla_benchmark_batched( mla_dims, vllm_config, device, + max_num_tokens=max_total_q, index_topk=index_topk if is_sparse else None, + kv_cache_dtype=kv_cache_dtype, ) # Verify the actual prefill backend matches what was requested @@ -942,6 +1040,7 @@ def _run_mla_benchmark_batched( mla_dims, device, indexer=indexer, + kv_cache_dtype=kv_cache_dtype, ) results.append(result) diff --git a/benchmarks/attention_benchmarks/runner.py b/benchmarks/attention_benchmarks/runner.py index 6af56e0e94f5..aa636cd9cb53 100644 --- a/benchmarks/attention_benchmarks/runner.py +++ b/benchmarks/attention_benchmarks/runner.py @@ -140,7 +140,7 @@ def _create_vllm_config( cache_config = CacheConfig( block_size=config.block_size, - cache_dtype="auto", + cache_dtype=config.kv_cache_dtype, ) cache_config.num_gpu_blocks = max_num_blocks cache_config.num_cpu_blocks = 0 @@ -215,7 +215,7 @@ def _create_backend_impl( num_kv_heads=config.num_kv_heads, alibi_slopes=None, sliding_window=None, - kv_cache_dtype="auto", + kv_cache_dtype=config.kv_cache_dtype, ) kv_cache_spec = FullAttentionSpec( @@ -288,12 +288,22 @@ def _create_input_tensors( total_q: int, device: torch.device, dtype: torch.dtype, + quantize_query: bool = False, ) -> tuple: - """Create Q, K, V input tensors for all layers.""" + """Create Q, K, V input tensors for all layers. + + When quantize_query is True, queries are cast to fp8 to match backends + that require query/key/value dtype consistency. + """ + q_dtype = dtype + if quantize_query: + from vllm.platforms import current_platform + + q_dtype = current_platform.fp8_dtype() q_list = [ torch.randn( total_q, config.num_q_heads, config.head_dim, device=device, dtype=dtype - ) + ).to(q_dtype) for _ in range(config.num_layers) ] k_list = [ @@ -344,10 +354,17 @@ def _create_kv_cache( # Compute inverse permutation to get back to logical view inv_order = [stride_order.index(i) for i in range(len(stride_order))] + # Use fp8 dtype for cache when requested. + cache_dtype = dtype + if config.kv_cache_dtype == "fp8": + from vllm.platforms import current_platform + + cache_dtype = current_platform.fp8_dtype() + cache_list = [] for _ in range(config.num_layers): # Allocate in physical layout order (contiguous in memory) - cache = torch.zeros(*physical_shape, device=device, dtype=dtype) + cache = torch.zeros(*physical_shape, device=device, dtype=cache_dtype) # Permute to logical view cache = cache.permute(*inv_order) cache_list.append(cache) @@ -392,6 +409,37 @@ def _run_single_benchmark( ) torch.accelerator.synchronize() + # Optionally capture a CUDA graph after warmup. + # Graph replay eliminates CPU launch overhead so timings reflect pure + # kernel time. + if config.use_cuda_graphs: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for i in range(config.num_layers): + impl.forward( + layer, + q_list[i], + k_list[i], + v_list[i], + cache_list[i], + attn_metadata, + output=out, + ) + benchmark_fn = graph.replay + else: + + def benchmark_fn(): + for i in range(config.num_layers): + impl.forward( + layer, + q_list[i], + k_list[i], + v_list[i], + cache_list[i], + attn_metadata, + output=out, + ) + # Benchmark times = [] for _ in range(config.repeats): @@ -399,16 +447,7 @@ def _run_single_benchmark( end = torch.cuda.Event(enable_timing=True) start.record() - for i in range(config.num_layers): - impl.forward( - layer, - q_list[i], - k_list[i], - v_list[i], - cache_list[i], - attn_metadata, - output=out, - ) + benchmark_fn() end.record() torch.accelerator.synchronize() @@ -502,8 +541,12 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult: common_attn_metadata=common_metadata, ) + # Only quantize queries when the impl supports it + quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr( + impl, "supports_quant_query_input", False + ) q_list, k_list, v_list = _create_input_tensors( - config, total_q, device, dtype + config, total_q, device, dtype, quantize_query=quantize_query ) cache_list = _create_kv_cache( diff --git a/csrc/ops.h b/csrc/ops.h index 921d6484d2d3..299650be73bf 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -295,10 +295,14 @@ void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a, std::vector cutlass_sparse_compress(torch::Tensor const& a); -void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input, - torch::Tensor& output_scale, - torch::Tensor const& input_scale, - bool is_sf_swizzled_layout); +std::tuple scaled_fp4_quant_func( + torch::Tensor const& input, torch::Tensor const& input_scale, + bool is_sf_swizzled_layout); + +void scaled_fp4_quant_out(torch::Tensor const& input, + torch::Tensor const& input_scale, + bool is_sf_swizzled_layout, torch::Tensor& output, + torch::Tensor& output_scale); void scaled_fp4_experts_quant( torch::Tensor& output, torch::Tensor& output_scale, diff --git a/csrc/quantization/fp4/nvfp4_quant_entry.cu b/csrc/quantization/fp4/nvfp4_quant_entry.cu index 650b9da8a499..8b5a1fd22cb7 100644 --- a/csrc/quantization/fp4/nvfp4_quant_entry.cu +++ b/csrc/quantization/fp4/nvfp4_quant_entry.cu @@ -16,6 +16,8 @@ #include +#include "nvfp4_utils.cuh" + #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, @@ -51,9 +53,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa( torch::Tensor const& output_scale_offset_by_experts); #endif -void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input, - torch::Tensor& output_sf, torch::Tensor const& input_sf, - bool is_sf_swizzled_layout) { +void scaled_fp4_quant_out(torch::Tensor const& input, + torch::Tensor const& input_sf, + bool is_sf_swizzled_layout, torch::Tensor& output, + torch::Tensor& output_sf) { #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \ (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120) return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf, @@ -62,6 +65,34 @@ void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input, TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel"); } +std::tuple scaled_fp4_quant_func( + torch::Tensor const& input, torch::Tensor const& input_sf, + bool is_sf_swizzled_layout) { + int64_t n = input.size(-1); + int64_t m = input.numel() / n; + auto device = input.device(); + + // Two fp4 values packed into a uint8 + auto output = torch::empty( + {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8)); + + torch::Tensor output_sf; + if (is_sf_swizzled_layout) { + auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n); + output_sf = torch::empty( + {sf_m, sf_n}, + torch::TensorOptions().device(device).dtype(torch::kInt32)); + } else { + output_sf = torch::empty( + {m, n / CVT_FP4_SF_VEC_SIZE}, + torch::TensorOptions().device(device).dtype(torch::kUInt8)); + } + + scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output, + output_sf); + return {output, output_sf}; +} + void scaled_fp4_experts_quant( torch::Tensor& output, torch::Tensor& output_scale, torch::Tensor const& input, torch::Tensor const& input_global_scale, diff --git a/csrc/quantization/fp4/nvfp4_utils.cuh b/csrc/quantization/fp4/nvfp4_utils.cuh index c1df1860c1a1..0c04f010888d 100644 --- a/csrc/quantization/fp4/nvfp4_utils.cuh +++ b/csrc/quantization/fp4/nvfp4_utils.cuh @@ -18,6 +18,7 @@ #include #include +#include #include "../../cuda_vec_utils.cuh" @@ -54,6 +55,18 @@ inline int computeEffectiveRows(int m) { return round_up(m, ROW_TILE); } +// Compute the shape of the swizzled SF output tensor. +// Returns (rounded_m, rounded_n / 4) where: +// rounded_m = round_up(m, 128) +// rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4) +inline std::pair computeSwizzledSFShape(int64_t m, + int64_t n) { + int64_t rounded_m = round_up(m, static_cast(128)); + int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE; + int64_t rounded_n = round_up(scale_n, static_cast(4)); + return {rounded_m, rounded_n / 4}; +} + // Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) { uint32_t val; diff --git a/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu index e178f252624f..723ca8142b82 100644 --- a/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu +++ b/csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu @@ -286,6 +286,15 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input, "Outer scale stride must be 1 when scales are not transposed"); } + int64_t hidden_size = input.size(-1); + TORCH_CHECK(hidden_size > 0 && hidden_size % group_size == 0, + "hidden_size must be a positive multiple of group_size"); + int64_t num_tokens = input.numel() / hidden_size; + int64_t num_groups = hidden_size / group_size; + TORCH_CHECK(scales.numel() >= num_tokens * num_groups, + "scales buffer too small: need ", num_tokens * num_groups, + " elements, got ", scales.numel()); + rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size, var_epsilon, scale_ub, residual, is_scale_transposed); diff --git a/csrc/sampler.cu b/csrc/sampler.cu index 30bfef33c0b0..2e76873c8f18 100644 --- a/csrc/sampler.cu +++ b/csrc/sampler.cu @@ -575,7 +575,7 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode( // The range of logits within the row. int rowStart = 0; int seq_len = seqLens[rowIdx / next_n]; - int rowEnd = seq_len - next_n + (rowIdx % next_n) + 1; + int rowEnd = max(0, seq_len - next_n + (rowIdx % next_n) + 1); // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) { diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index d98e987d92a2..aadc9fe33753 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -564,10 +564,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Compute NVFP4 block quantized tensor. ops.def( - "scaled_fp4_quant(Tensor! output, Tensor input," - " Tensor! output_scale, Tensor input_scale, bool " - "is_sf_swizzled_layout) -> ()"); - ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant); + "scaled_fp4_quant(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout) -> (Tensor, Tensor)"); + ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func); + + // Out variant + // TODO: Add {at::Tag::out_variant} tag and update all call sites + // to use the functional variant once vLLM upgrades PyTorch. + // See pytorch/pytorch#176117. + ops.def( + "scaled_fp4_quant.out(Tensor input," + " Tensor input_scale, bool " + "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) " + "-> ()"); + ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out); // Compute NVFP4 experts quantization. ops.def( diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index a8d2fd687fff..7c60a136f790 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -164,18 +164,18 @@ Priority is **1 = highest** (tried first). | Backend | Version | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------- | ------ | --------- | ----------- | ---------- | ---- | --------- | --- | --------------- | ------------ | | `CPU_ATTN` | | fp16, bf16, fp32 | `auto` | Any | 32, 64, 80, 96, 112, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | All | N/A | -| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | -| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x | -| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 | -| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x | -| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 | +| `FLASHINFER` | Native† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ❌ | ❌ | ✅ | Decoder | 7.x-9.x | +| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x | +| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 | +| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x | +| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 | | `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any | -| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any | -| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A | +| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any | +| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder, Enc-Dec | N/A | | `ROCM_AITER_UNIFIED_ATTN` | | fp16, bf16 | `auto` | %16 | Any | ✅ | ✅ | ❌ | All | N/A | -| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A | -| `TREE_ATTN` | | fp16, bf16 | `auto` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | -| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any | +| `ROCM_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | 32, 64, 80, 96, 128, 160, 192, 224, 256 | ✅ | ✅ | ❌ | All | N/A | +| `TREE_ATTN` | | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | 32, 64, 96, 128, 160, 192, 224, 256 | ❌ | ❌ | ❌ | Decoder | Any | +| `TRITON_ATTN` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ✅ | ❌ | All | Any | > **†** FlashInfer uses TRTLLM attention on Blackwell (SM100), which supports sinks. Disable via `--attention-config.use_trtllm_attention=0`. > @@ -204,14 +204,14 @@ configuration. | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. | | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ | -| `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | -| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | -| `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | +| `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x | +| `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x | +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x | +| `FLASHMLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x | | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x | -| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | -| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | +| `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x | +| `ROCM_AITER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 1 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | +| `ROCM_AITER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | 1 | Any | ❌ | ✅ | ❌ | ❌ | Decoder | N/A | | `ROCM_AITER_TRITON_MLA` | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ❌ | ❌ | Decoder | N/A | -| `TRITON_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | -| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any | +| `TRITON_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | Any | +| `XPU_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16` | Any | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | Any | diff --git a/docs/design/torch_compile_multimodal.md b/docs/design/torch_compile_multimodal.md index 4abf1d08c517..c46bfa8325bb 100644 --- a/docs/design/torch_compile_multimodal.md +++ b/docs/design/torch_compile_multimodal.md @@ -34,9 +34,6 @@ relies on caching artifacts to reduce start time, we must properly propagate the with the LLM text-backbone, or other instances of the same artifact (as is the case with vision block). `is_encoder=True` is also needed for encoder components (see Compile Range Integration). -3. `with set_forward_context` context manager should be used around the nn.Module's forward call. This will properly forward the vllm_config which is needed -for torch.compile integration. - ### CompilationConfig With the exception of `compile_mm_encoder: true`, the multimodal encoder will inherit from the same compilation config as the text LLM. We may extend diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index fe95735b91b0..cea1175413fe 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -107,6 +107,27 @@ vLLM supports the `tool_choice='none'` option in the chat completion API. When t !!! note When tools are specified in the request, vLLM includes tool definitions in the prompt by default, regardless of the `tool_choice` setting. To exclude tool definitions when `tool_choice='none'`, use the `--exclude-tools-when-tool-choice-none` option. +## Constrained Decoding Behavior + +Whether vLLM enforces the tool parameter schema during generation depends on the `tool_choice` mode: + +| `tool_choice` value | Schema-constrained decoding | Behavior | +| --- | --- | --- | +| Named function | Yes (via structured outputs backend) | Arguments are guaranteed to be valid JSON conforming to the function's parameter schema. | +| `"required"` | Yes (via structured outputs backend) | Same as named function. The model must produce at least one tool call. | +| `"auto"` | No | The model generates freely. A tool-call parser extracts tool calls from the raw text. Arguments may be malformed or not match the schema. | +| `"none"` | N/A | No tool calls are produced. | + +When schema conformance matters, prefer `tool_choice="required"` or named function calling over `"auto"`. + +### Strict Mode (`strict` parameter) + +The [OpenAI API](https://platform.openai.com/docs/guides/function-calling#strict-mode) supports a `strict` field on function definitions. When set to `true`, OpenAI uses constrained decoding to guarantee that tool-call arguments match the function schema, even in `tool_choice="auto"` mode. + +vLLM **does not implement** `strict` mode today. The `strict` field is accepted in requests (to avoid breaking clients that set it), but it has no effect on decoding behavior. In auto mode, argument validity depends entirely on the model's output quality and the parser's extraction logic. + +Tracking issues: [#15526](https://github.com/vllm-project/vllm/issues/15526), [#16313](https://github.com/vllm-project/vllm/issues/16313). + ## Automatic Function Calling To enable this feature, you should set the following flags: @@ -124,6 +145,9 @@ from HuggingFace; and you can find an example of this in a `tokenizer_config.jso If your favorite tool-calling model is not supported, please feel free to contribute a parser & tool use chat template! +!!! note + With `tool_choice="auto"`, tool-call arguments are extracted from the model's raw text output by the selected parser. No schema-level constraint is applied during decoding, so arguments may occasionally be malformed or violate the function's parameter schema. See [Constrained Decoding Behavior](#constrained-decoding-behavior) for details. + ### Hermes Models (`hermes`) All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported. @@ -219,7 +243,7 @@ Supported models: * `ibm-granite/granite-4.0-h-small` and other Granite 4.0 models - Recommended flags: `--tool-call-parser hermes` + Recommended flags: `--tool-call-parser granite4` * `ibm-granite/granite-3.0-8b-instruct` diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index fe50081e5ce7..92e7402c0537 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -179,7 +179,7 @@ def ops_in_model_after(self): def ops_in_model_before(self): return [ torch.ops.vllm.all_reduce.default, - torch.ops._C.scaled_fp4_quant.default, + torch.ops._C.scaled_fp4_quant.out, ] diff --git a/tests/entrypoints/openai/test_gptoss_structural_tags_integration.py b/tests/entrypoints/openai/test_gptoss_structural_tags_integration.py deleted file mode 100644 index e9d33ba9bbb8..000000000000 --- a/tests/entrypoints/openai/test_gptoss_structural_tags_integration.py +++ /dev/null @@ -1,279 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Integration tests for GPT-OSS structural tags functionality (PR #25515).""" - -import json -from unittest.mock import Mock - -import pytest - -from vllm.entrypoints.mcp.tool_server import ToolServer -from vllm.reasoning.gptoss_reasoning_parser import ( - GptOssReasoningParser, -) -from vllm.sampling_params import StructuredOutputsParams - - -class TestGptOssStructuralTagsIntegration: - """Integration tests for structural tags in GPT-OSS tool calls.""" - - @pytest.fixture - def mock_tokenizer(self): - """Create a mock tokenizer.""" - tokenizer = Mock() - tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5]) - tokenizer.get_vocab = Mock(return_value={"<|end|>": 6}) - return tokenizer - - @pytest.fixture - def gptoss_parser(self, mock_tokenizer): - """Create a real GptOssReasoningParser instance.""" - return GptOssReasoningParser(mock_tokenizer) - - @pytest.fixture - def tool_server_with_python(self): - """Create a tool server with Python tool enabled.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python") - return tool_server - - @pytest.fixture - def tool_server_empty(self): - """Create a tool server with no tools.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(return_value=False) - return tool_server - - def test_end_to_end_no_tools(self, gptoss_parser): - """Test end-to-end flow when no tools are available.""" - # Test the parser directly - result = gptoss_parser.prepare_structured_tag(None, None) - parsed_result = json.loads(result) - - # Verify basic structure - assert parsed_result["type"] == "structural_tag" - assert parsed_result["format"]["type"] == "triggered_tags" - assert len(parsed_result["format"]["tags"]) == 1 - - # Verify only analysis channel is allowed - analysis_tag = parsed_result["format"]["tags"][0] - assert analysis_tag["begin"] == "<|channel|>analysis<|message|>" - assert analysis_tag["content"]["type"] == "any_text" - assert analysis_tag["end"] == "<|end|>" - - # Verify triggers - assert parsed_result["format"]["triggers"] == ["<|channel|>analysis"] - assert parsed_result["format"]["stop_after_first"] is False - - def test_end_to_end_with_python_tool(self, gptoss_parser, tool_server_with_python): - """Test end-to-end flow with Python tool enabled.""" - result = gptoss_parser.prepare_structured_tag(None, tool_server_with_python) - parsed_result = json.loads(result) - - # Should have analysis tag + 2 python tags - assert len(parsed_result["format"]["tags"]) == 3 - - # Verify all expected tags are present - tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]] - expected_begins = [ - "<|channel|>analysis<|message|>", - "<|channel|>commentary to=python", - "<|channel|>analysis to=python", - ] - - for expected in expected_begins: - assert expected in tag_begins - - # Verify triggers include commentary - assert "<|channel|>analysis" in parsed_result["format"]["triggers"] - assert "<|channel|>commentary to=" in parsed_result["format"]["triggers"] - - def test_structured_outputs_params_integration( - self, gptoss_parser, tool_server_with_python - ): - """Test integration with StructuredOutputsParams.""" - # Generate structural tag - structural_tag = gptoss_parser.prepare_structured_tag( - None, tool_server_with_python - ) - - # Create StructuredOutputsParams - params = StructuredOutputsParams(structural_tag=structural_tag) - - # Verify the tag is properly stored and accessible - assert params.structural_tag == structural_tag - - # Verify the tag is valid JSON - parsed_tag = json.loads(params.structural_tag) - assert parsed_tag["type"] == "structural_tag" - - @pytest.mark.parametrize( - "browser, python, container, expected_tags", - [ - # No tools - (False, False, False, 1), - # Single tool - (True, False, False, 3), - # Multiple tools - (True, True, False, 5), - # All tools - (True, True, True, 7), - ], - ) - def test_tool_server_interaction_flow( - self, gptoss_parser, browser, python, container, expected_tags - ): - """Test the complete tool server interaction flow.""" - - # Create a mock ToolServer - tool_server = Mock(spec=ToolServer) - - # Simulate tool availability based on parameters - tool_server.has_tool = Mock( - side_effect=lambda tool: { - "browser": browser, - "python": python, - "container": container, - }.get(tool, False) - ) - - # Run the parser and verify results - result = gptoss_parser.prepare_structured_tag(None, tool_server) - parsed_result = json.loads(result) - - # Validate number of tags - assert len(parsed_result["format"]["tags"]) == expected_tags - - # Verify tool-specific tags exist for enabled tools - tag_begins = [tag["begin"] for tag in parsed_result["format"]["tags"]] - for tool, enabled in { - "browser": browser, - "python": python, - "container": container, - }.items(): - if enabled: - assert f"<|channel|>commentary to={tool}" in tag_begins - assert f"<|channel|>analysis to={tool}" in tag_begins - - def test_original_tag_preservation(self, gptoss_parser, tool_server_with_python): - """Test that original tags are preserved when provided.""" - original_tag = '{"type": "custom_tag", "data": "preserved"}' - - result = gptoss_parser.prepare_structured_tag( - original_tag, tool_server_with_python - ) - - # Should return original tag unchanged - assert result == original_tag - - @pytest.mark.parametrize( - "tools", - [ - [], - ["browser"], - ["python"], - ["container"], - ["browser", "python"], - ["browser", "container"], - ["python", "container"], - ["browser", "python", "container"], - ], - ) - def test_json_validity_comprehensive(self, gptoss_parser, tools): - """Test JSON validity across all possible tool combinations.""" - - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools) - - result = gptoss_parser.prepare_structured_tag(None, tool_server) - - # Should be valid JSON - parsed_result = json.loads(result) - - # Should have correct structure - assert parsed_result["type"] == "structural_tag" - assert "format" in parsed_result - assert "tags" in parsed_result["format"] - assert "triggers" in parsed_result["format"] - - # Tag count should be: 1 (analysis) + 2 * len(tools) - expected_tag_count = 1 + (2 * len(tools)) - assert len(parsed_result["format"]["tags"]) == expected_tag_count - - def test_error_handling_invalid_tool_server(self, gptoss_parser): - """Test error handling with invalid tool server.""" - # Tool server that raises exceptions - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=Exception("Tool server error")) - - # Should handle gracefully and still return a valid tag - with pytest.raises(Exception, match="Tool server error"): - gptoss_parser.prepare_structured_tag(None, tool_server) - - def test_concurrent_requests_isolation(self, gptoss_parser): - """Test that concurrent requests don't interfere with each other.""" - # Simulate concurrent requests with different tool servers - tool_server_1 = Mock(spec=ToolServer) - tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python") - - tool_server_2 = Mock(spec=ToolServer) - tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser") - - # Generate tags concurrently - result_1 = gptoss_parser.prepare_structured_tag(None, tool_server_1) - result_2 = gptoss_parser.prepare_structured_tag(None, tool_server_2) - - # Parse results - parsed_1 = json.loads(result_1) - parsed_2 = json.loads(result_2) - - # Verify they have different tool configurations - tags_1 = [tag["begin"] for tag in parsed_1["format"]["tags"]] - tags_2 = [tag["begin"] for tag in parsed_2["format"]["tags"]] - - # Result 1 should have python tags - assert "<|channel|>commentary to=python" in tags_1 - assert "<|channel|>commentary to=browser" not in tags_1 - - # Result 2 should have browser tags - assert "<|channel|>commentary to=browser" in tags_2 - assert "<|channel|>commentary to=python" not in tags_2 - - def test_tag_format_consistency(self, gptoss_parser): - """Test that all generated tags follow consistent format.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock( - side_effect=lambda tool: tool in ["python", "browser"] - ) - - result = gptoss_parser.prepare_structured_tag(None, tool_server) - parsed_result = json.loads(result) - - # Verify all tags have required fields - for tag in parsed_result["format"]["tags"]: - assert "begin" in tag - assert "content" in tag - assert "end" in tag - assert tag["content"]["type"] == "any_text" - assert tag["end"] == "<|end|>" - - # Verify begin format - assert tag["begin"].startswith("<|channel|>") - - def test_trigger_configuration(self, gptoss_parser): - """Test trigger configuration for different tool setups.""" - # Test with no tools - result_no_tools = gptoss_parser.prepare_structured_tag(None, None) - parsed_no_tools = json.loads(result_no_tools) - assert parsed_no_tools["format"]["triggers"] == ["<|channel|>analysis"] - - # Test with tools - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=lambda tool: tool == "python") - - result_with_tools = gptoss_parser.prepare_structured_tag(None, tool_server) - parsed_with_tools = json.loads(result_with_tools) - - expected_triggers = ["<|channel|>analysis", "<|channel|>commentary to="] - assert set(parsed_with_tools["format"]["triggers"]) == set(expected_triggers) diff --git a/tests/entrypoints/openai/tool_parsers/test_granite4_tool_parser.py b/tests/entrypoints/openai/tool_parsers/test_granite4_tool_parser.py new file mode 100644 index 000000000000..27e7a8c5dabf --- /dev/null +++ b/tests/entrypoints/openai/tool_parsers/test_granite4_tool_parser.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +import random +from typing import Any + +import openai +import pytest +from transformers import AutoTokenizer + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest +from vllm.entrypoints.openai.engine.protocol import ( + DeltaMessage, +) +from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser + +from ....utils import RemoteOpenAIServer + +MODEL = "ibm-granite/granite-4.0-h-tiny" + + +@pytest.fixture(scope="module") +def server(): + model = MODEL + args_for_model = [ + "--enforce-eager", + "--enable-auto-tool-choice", + "--tool-call-parser", + "granite4", + "--tokenizer", + "ibm-granite/granite-4.0-h-tiny", + "--max-model-len", + "4096", + "--max-num-seqs", + "2", + ] + with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server: + yield server + + +def create_complex_input(create_string_args: bool): + coord_arg: dict | str = { + "coordinates": [[23.54, 43.1], [-12.2, 54.3], [4, 5]], + "coordinate_type": "latlong", + } + if create_string_args: + # test granite behavior + coord_arg = json.dumps(coord_arg) + return [ + {"name": "find_bbox", "arguments": coord_arg}, + { + "name": "get_stock_price", + "arguments": { + "symbol": "AAPL", + "start_date": "2021-01-01", + "end_date": "2021-12-31", + }, + }, + {"name": "find_bbox", "arguments": coord_arg}, + ] + + +def random_chunks(s: str, min_len: int, max_len: int): + chunks = [] + i = 0 + n = len(s) + + while i < n: + size = random.randint(min_len, max_len) + chunks.append(s[i : i + size]) + i += size + + return chunks + + +@pytest.fixture(scope="module") +def tokenizer(): + return AutoTokenizer.from_pretrained(MODEL) + + +# create a variety of input chunk sizes +@pytest.mark.parametrize( + "min_chunk, max_chunk", + [ + (1, 1), + (1, 2), + (5, 7), + (6, 20), + ], +) +def test_tool_call_parser_complex(min_chunk: int, max_chunk: int, tokenizer): + input_dicts = create_complex_input(True) + + formatted_tcs = [ + " " + json.dumps(call) + " " for call in input_dicts + ] + + text_messages = [ + "Here goes the bbox call: \n", + " Now the stock price call: \n ", + " Now another bbox call: \n ", + " See? I'm a helpful assistant.", + ] + + test_input = ( + text_messages[0] + + formatted_tcs[0] + + text_messages[1] + + formatted_tcs[1] + + text_messages[2] + + formatted_tcs[2] + + text_messages[3] + ) + + any_chat_request = ChatCompletionRequest( + seed=42, + model=MODEL, + messages=[], + ) + + parser = Granite4ToolParser(tokenizer=tokenizer) + + delta_messages = list[DeltaMessage]() + for text in random_chunks(test_input, min_chunk, max_chunk): + delta = parser.extract_tool_calls_streaming( + previous_text="", + current_text="", + delta_text=text, + previous_token_ids=[], + current_token_ids=[], + delta_token_ids=[], + request=any_chat_request, + ) + if delta is not None: + delta_messages.append(delta) + + content = "" + tool_calls = list[dict[str, Any]]() + + current_name = "__start__" + current_args = "" + + for msg in delta_messages: + if msg.content: + content += msg.content + for tool_call in msg.tool_calls: + if delta_func := tool_call.function: + if delta_func.name is not None: + if current_name == "__start__": + current_name = delta_func.name + + if delta_func.name != current_name: + tool_calls.append( + { + "name": current_name, + "arguments": json.loads(current_args), + } + ) + current_name = delta_func.name + current_args = "" + + if delta_func.arguments: + current_args += delta_func.arguments + + if current_name != "__start__": + tool_calls.append({"name": current_name, "arguments": json.loads(current_args)}) + + assert content == "".join(text_messages) + assert tool_calls == create_complex_input(False) + + +tools = [ + { + "type": "function", + "function": { + "name": "get_acme_region_name_for_transaction_id", + "description": "Returns ACME transaction/transaction ID information" + " including ACME regions\n\nArgs:\n start_time " + "(str): Start date and time in datetime format " + '"%Y-%m-%dT%H:%M:%S.%f"\n end_time (str): End ' + "date and time in datetime format " + '"%Y-%m-%dT%H:%M:%S.%f"\n size (int, optional): ' + "Number of ACME Transaction IDs to return\n " + "order (str, optional): Sort by most run " + "transaction IDs. The value can be 'asc' for " + "ascending or 'desc' for descending\n " + "transaction_id (str, optional): ACME Transaction " + "ID to filter on\n acme_region (str, optional): " + "ACME Region to filter on\nReturns:\n - A " + "dictionary containing a list of ACME transaction " + "ids and the ACME regions they run in:\n {\n" + ' "Number of transaction IDs" : int,\n' + ' "Total transaction IDs available": int' + ',\n "ACME Transaction IDs": [\n ' + ' {\n "Transaction ID": ' + 'str,\n "Number of runs": int,\n' + ' "ACME Regions": [str],\n ' + " },\n ...\n ]," + '\n "Start time" : datetime,\n ' + ' "End time" : datetime,\n ' + ' "Order" : str\n }\n ' + " - If no ACME region found for transaction id, " + 'returns:\n {"Success": "No ACME region ' + 'found for transaction id."}\n - If an error ' + 'occurs, returns:\n {"Error": "{exception' + ' message}"}', + "parameters": { + "properties": { + "start_time": {}, + "end_time": {}, + "size": {"default": 500}, + "order": {"default": "desc"}, + "transaction_id": {"default": None}, + "acme_region": {"default": None}, + }, + "required": ["start_time", "end_time"], + "type": "object", + }, + }, + } +] + +tools2 = [ + { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": { + "description": "The city and state, e.g. San Francisco, CA", + "type": "string", + } + }, + "required": ["location"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_stock_price", + "description": "Retrieves the current stock price for a given " + "ticker symbol. The ticker symbol must be a valid " + "symbol for a publicly traded company on a major US" + " stock exchange like NYSE or NASDAQ. The tool will" + " return the latest trade price in USD. It should " + "be used when the user asks about the current or " + "most recent price of a specific stock. It will not" + " provide any other information about the stock or" + " company.", + "parameters": { + "type": "object", + "properties": { + "ticker": { + "description": "The stock ticker symbol, e.g." + " AAPL for Apple Inc.", + "type": "string", + } + }, + }, + }, + }, +] + +messages = [ + { + "content": "\n\nSystem: You are a helpful, precise, and methodical AI" + " assistant that uses tool outputs provided inline.\nAlways" + " assume the current datetime is 2026-01-29T13:59:09.238901" + "+00:00.\n\nIf you receive a ToolMessage with `tool_call_id" + '` equal to "get_time_range" (or "time_range_tool"), you ' + "MUST:\n 1. Parse that JSON and use the values `start` and" + " `end` directly when calling other tools.\n 2. Do not " + "re-call or re-compute the time range.\n 3. Pass resolved " + "values (ISO strings) as arguments to any subsequent tool " + "(do not pass function metadata or placeholders).\n 4. If " + "a tool requires datetime objects rather than strings, " + "convert the ISO strings into language-native datetime " + "objects before invoking.\n\nAlways return fully resolved " + "arguments in correct types (e.g., ISO datetime strings or" + " datetime objects) and never include placeholders like " + '"".\n\n', + "role": "system", + }, + { + "content": "What are the transaction IDs that ran in the" + " ACME region A9345 over the last two months?", + "role": "user", + }, + { + "content": '["2026-01-26T09: 51: 55.467722Z", "2026-01-27T09: 51: 55.467722Z"]', + "role": "tool", + "tool_call_id": "time_range_tool", + }, +] +messages2 = [{"role": "user", "content": "What's stock price for IBM?"}] + +messages3 = [{"role": "user", "content": "What's the current weather in New York?"}] + + +def get_args(client: openai.OpenAI, _tools, _messages, _stop): + response = client.chat.completions.create( + model=MODEL, + messages=_messages, + temperature=0, + tools=_tools, + max_tokens=200, + stop=_stop, + tool_choice="auto", + ) + + return response.choices[0].message.tool_calls[0].function.arguments + + +async def get_args_streaming( + async_client: openai.AsyncOpenAI, _tools, _messages, _stop +): + stream = await async_client.chat.completions.create( + model=MODEL, + messages=_messages, + temperature=0, + tools=_tools, + max_tokens=200, + stop=_stop, + tool_choice="auto", + stream=True, + ) + full_call = [] + async for chunk in stream: + tc = chunk.choices[0].delta.tool_calls + if tc and tc[0].function.arguments: + full_call.append(tc[0].function.arguments) + return "".join(full_call) + + +async def run_scenario(server: RemoteOpenAIServer, _tools, _messages, _stop): + non_streaming = get_args(server.get_client(), _tools, _messages, _stop) + json.loads(non_streaming) # verify that it is json loadable + streaming = await get_args_streaming( + server.get_async_client(), _tools, _messages, _stop + ) + json.loads(streaming) + assert non_streaming == streaming, f"{non_streaming=}, {streaming=}" + + +@pytest.mark.asyncio +async def test_stop_sequence_interference(server: RemoteOpenAIServer): + print("Testing scenario 1") + await run_scenario(server, tools, messages, "veroniqueprattyushveroniqueprattyush") + + print("Testing scenario 2") + await run_scenario( + server, tools2, messages2, "veroniqueprattyushveroniqueprattyush" + ) + + print("Testing scenario 3") + await run_scenario(server, tools2, messages3, "prattyush") diff --git a/tests/entrypoints/openai/tool_parsers/test_hermes_tool_parser.py b/tests/entrypoints/openai/tool_parsers/test_hermes_tool_parser.py index 626d845e1b44..be910fbb1a41 100644 --- a/tests/entrypoints/openai/tool_parsers/test_hermes_tool_parser.py +++ b/tests/entrypoints/openai/tool_parsers/test_hermes_tool_parser.py @@ -3,29 +3,22 @@ import json +import openai import pytest +import pytest_asyncio +from huggingface_hub import snapshot_download +from typing_extensions import TypedDict from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.tokenizers import TokenizerLike +from vllm.tool_parsers.abstract_tool_parser import ToolParser +from vllm.tool_parsers.granite4_tool_parser import Granite4ToolParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser from ....utils import RemoteOpenAIServer -MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct" LORA_MODEL = "minpeter/LoRA-Llama-3.2-1B-tool-vllm-ci" -SERVER_ARGS = [ - "--enforce-eager", - "--enable-auto-tool-choice", - "--tool-call-parser", - "hermes", - "--enable-lora", - "--lora-modules", - f"{LORA_MODEL}={LORA_MODEL}", - "--tokenizer", - f"{LORA_MODEL}", -] - TOOLS = [ { "type": "function", @@ -50,6 +43,75 @@ } ] + +class ServerConfig(TypedDict, total=False): + model: str + arguments: list[str] + model_arg: str + tool_parser: ToolParser + + +CONFIGS: dict[str, ServerConfig] = { + "llama": { + "model": "meta-llama/Llama-3.2-1B-Instruct", + "arguments": [ + "--enforce-eager", + "--enable-auto-tool-choice", + "--tool-call-parser", + "hermes", + "--enable-lora", + "--lora-modules", + f"{LORA_MODEL}={LORA_MODEL}", + "--tokenizer", + f"{LORA_MODEL}", + ], + "model_arg": LORA_MODEL, + "tool_parser": Hermes2ProToolParser, + }, + "granite4": { + "model": "ibm-granite/granite-4.0-h-tiny", + "arguments": [ + "--enforce-eager", + "--enable-auto-tool-choice", + "--tool-call-parser", + "granite4", + "--tokenizer", + "ibm-granite/granite-4.0-h-tiny", + "--max-model-len", + "4096", + "--max-num-seqs", + "2", + ], + "model_arg": "ibm-granite/granite-4.0-h-tiny", + "tool_parser": Granite4ToolParser, + }, +} + + +# for each server config, download the model and return the config +@pytest.fixture(scope="session", params=CONFIGS.keys()) +def server_config(request): + config = CONFIGS[request.param] + + # download model and tokenizer using transformers + snapshot_download(config["model"]) + yield CONFIGS[request.param] + + +@pytest.fixture(scope="module") +def server(request, server_config: ServerConfig): + model = server_config["model"] + args_for_model = server_config["arguments"] + with RemoteOpenAIServer(model, args_for_model, max_wait_seconds=480) as server: + yield server + + +@pytest_asyncio.fixture +async def client(server: RemoteOpenAIServer): + async with server.get_async_client() as async_client: + yield async_client + + PRODUCT_TOOLS = [ { "type": "function", @@ -87,186 +149,182 @@ @pytest.mark.asyncio -async def test_non_streaming_tool_call(): +async def test_non_streaming_tool_call( + client: openai.AsyncOpenAI, server_config: ServerConfig +): """Test tool call in non-streaming mode.""" - with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server: - client = server.get_async_client() - - response = await client.chat.completions.create( - model=LORA_MODEL, - messages=MESSAGES, - tools=TOOLS, - tool_choice="auto", - temperature=0.0, - ) - assert response.choices - choice = response.choices[0] - message = choice.message + response = await client.chat.completions.create( + model=server_config["model_arg"], + messages=MESSAGES, + tools=TOOLS, + tool_choice="auto", + temperature=0.0, + ) + + assert response.choices + choice = response.choices[0] + message = choice.message - assert choice.finish_reason == "tool_calls" - assert message.tool_calls is not None + assert choice.finish_reason == "tool_calls" + assert message.tool_calls is not None - tool_call = message.tool_calls[0] - assert tool_call.type == "function" - assert tool_call.function.name == "get_current_weather" + tool_call = message.tool_calls[0] + assert tool_call.type == "function" + assert tool_call.function.name == "get_current_weather" - arguments = json.loads(tool_call.function.arguments) - assert "location" in arguments - assert "Boston" in arguments["location"] - print("\n[Non-Streaming Test Passed]") - print(f"Tool Call: {tool_call.function.name}") - print(f"Arguments: {arguments}") + arguments = json.loads(tool_call.function.arguments) + assert "location" in arguments + assert "Boston" in arguments["location"] + print("\n[Non-Streaming Test Passed]") + print(f"Tool Call: {tool_call.function.name}") + print(f"Arguments: {arguments}") @pytest.mark.asyncio -async def test_streaming_tool_call(): +async def test_streaming_tool_call( + client: openai.AsyncOpenAI, server_config: ServerConfig +): """Test tool call in streaming mode.""" - with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server: - client = server.get_async_client() - - stream = await client.chat.completions.create( - model=LORA_MODEL, - messages=MESSAGES, - tools=TOOLS, - tool_choice="auto", - temperature=0.0, - stream=True, - ) - tool_call_chunks = {} - async for chunk in stream: - if not chunk.choices: - continue + stream = await client.chat.completions.create( + model=server_config["model_arg"], + messages=MESSAGES, + tools=TOOLS, + tool_choice="auto", + temperature=0.0, + stream=True, + ) + + tool_call_chunks = {} + async for chunk in stream: + if not chunk.choices: + continue - delta = chunk.choices[0].delta - if not delta or not delta.tool_calls: - continue + delta = chunk.choices[0].delta + if not delta or not delta.tool_calls: + continue - for tool_chunk in delta.tool_calls: - index = tool_chunk.index - if index not in tool_call_chunks: - tool_call_chunks[index] = {"name": "", "arguments": ""} + for tool_chunk in delta.tool_calls: + index = tool_chunk.index + if index not in tool_call_chunks: + tool_call_chunks[index] = {"name": "", "arguments": ""} - if tool_chunk.function.name: - tool_call_chunks[index]["name"] += tool_chunk.function.name - if tool_chunk.function.arguments: - tool_call_chunks[index]["arguments"] += ( - tool_chunk.function.arguments - ) + if tool_chunk.function.name: + tool_call_chunks[index]["name"] += tool_chunk.function.name + if tool_chunk.function.arguments: + tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments - assert len(tool_call_chunks) == 1 - reconstructed_tool_call = tool_call_chunks[0] + assert len(tool_call_chunks) == 1 + reconstructed_tool_call = tool_call_chunks[0] - assert reconstructed_tool_call["name"] == "get_current_weather" + assert reconstructed_tool_call["name"] == "get_current_weather" - arguments = json.loads(reconstructed_tool_call["arguments"]) - assert "location" in arguments - assert "Boston" in arguments["location"] - print("\n[Streaming Test Passed]") - print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}") - print(f"Reconstructed Arguments: {arguments}") + arguments = json.loads(reconstructed_tool_call["arguments"]) + assert "location" in arguments + assert "Boston" in arguments["location"] + print("\n[Streaming Test Passed]") + print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}") + print(f"Reconstructed Arguments: {arguments}") @pytest.mark.asyncio -async def test_non_streaming_product_tool_call(): +async def test_non_streaming_product_tool_call( + client: openai.AsyncOpenAI, server_config: ServerConfig +): """Test tool call integer and boolean parameters in non-streaming mode.""" - with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server: - client = server.get_async_client() - - response = await client.chat.completions.create( - model=LORA_MODEL, - messages=PRODUCT_MESSAGES, - tools=PRODUCT_TOOLS, - tool_choice="auto", - temperature=0.66, - ) - assert response.choices - choice = response.choices[0] - message = choice.message + response = await client.chat.completions.create( + model=server_config["model_arg"], + messages=PRODUCT_MESSAGES, + tools=PRODUCT_TOOLS, + tool_choice="auto", + temperature=0.66, + ) + + assert response.choices + choice = response.choices[0] + message = choice.message - assert choice.finish_reason == "tool_calls" - assert message.tool_calls is not None + assert choice.finish_reason == "tool_calls" + assert message.tool_calls is not None - tool_call = message.tool_calls[0] - assert tool_call.type == "function" - assert tool_call.function.name == "get_product_info" + tool_call = message.tool_calls[0] + assert tool_call.type == "function" + assert tool_call.function.name == "get_product_info" - arguments = json.loads(tool_call.function.arguments) - assert "product_id" in arguments - assert "inserted" in arguments + arguments = json.loads(tool_call.function.arguments) + assert "product_id" in arguments + assert "inserted" in arguments - product_id = arguments.get("product_id") - inserted = arguments.get("inserted") + product_id = arguments.get("product_id") + inserted = arguments.get("inserted") - assert isinstance(product_id, int) - assert product_id == 7355608 - assert isinstance(inserted, bool) - assert inserted is True + assert isinstance(product_id, int) + assert product_id == 7355608 + assert isinstance(inserted, bool) + assert inserted is True - print("\n[Non-Streaming Product Test Passed]") - print(f"Tool Call: {tool_call.function.name}") - print(f"Arguments: {arguments}") + print("\n[Non-Streaming Product Test Passed]") + print(f"Tool Call: {tool_call.function.name}") + print(f"Arguments: {arguments}") @pytest.mark.asyncio -async def test_streaming_product_tool_call(): +async def test_streaming_product_tool_call( + client: openai.AsyncOpenAI, server_config: ServerConfig +): """Test tool call integer and boolean parameters in streaming mode.""" - with RemoteOpenAIServer(MODEL_NAME, SERVER_ARGS) as server: - client = server.get_async_client() - - stream = await client.chat.completions.create( - model=LORA_MODEL, - messages=PRODUCT_MESSAGES, - tools=PRODUCT_TOOLS, - tool_choice="auto", - temperature=0.66, - stream=True, - ) - tool_call_chunks = {} - async for chunk in stream: - if not chunk.choices: - continue + stream = await client.chat.completions.create( + model=server_config["model_arg"], + messages=PRODUCT_MESSAGES, + tools=PRODUCT_TOOLS, + tool_choice="auto", + temperature=0.66, + stream=True, + ) + + tool_call_chunks = {} + async for chunk in stream: + if not chunk.choices: + continue - delta = chunk.choices[0].delta - if not delta or not delta.tool_calls: - continue + delta = chunk.choices[0].delta + if not delta or not delta.tool_calls: + continue - for tool_chunk in delta.tool_calls: - index = tool_chunk.index - if index not in tool_call_chunks: - tool_call_chunks[index] = {"name": "", "arguments": ""} + for tool_chunk in delta.tool_calls: + index = tool_chunk.index + if index not in tool_call_chunks: + tool_call_chunks[index] = {"name": "", "arguments": ""} - if tool_chunk.function.name: - tool_call_chunks[index]["name"] += tool_chunk.function.name - if tool_chunk.function.arguments: - tool_call_chunks[index]["arguments"] += ( - tool_chunk.function.arguments - ) + if tool_chunk.function.name: + tool_call_chunks[index]["name"] += tool_chunk.function.name + if tool_chunk.function.arguments: + tool_call_chunks[index]["arguments"] += tool_chunk.function.arguments - assert len(tool_call_chunks) == 1 - reconstructed_tool_call = tool_call_chunks[0] + assert len(tool_call_chunks) == 1 + reconstructed_tool_call = tool_call_chunks[0] - assert reconstructed_tool_call["name"] == "get_product_info" + assert reconstructed_tool_call["name"] == "get_product_info" - arguments = json.loads(reconstructed_tool_call["arguments"]) - assert "product_id" in arguments - assert "inserted" in arguments + arguments = json.loads(reconstructed_tool_call["arguments"]) + assert "product_id" in arguments + assert "inserted" in arguments - # Handle type coercion for streaming test as well - product_id = arguments.get("product_id") - inserted = arguments.get("inserted") + # Handle type coercion for streaming test as well + product_id = arguments.get("product_id") + inserted = arguments.get("inserted") - assert isinstance(product_id, int) - assert product_id == 7355608 - assert isinstance(inserted, bool) - assert inserted is True + assert isinstance(product_id, int) + assert product_id == 7355608 + assert isinstance(inserted, bool) + assert inserted is True - print("\n[Streaming Product Test Passed]") - print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}") - print(f"Reconstructed Arguments: {arguments}") + print("\n[Streaming Product Test Passed]") + print(f"Reconstructed Tool Call: {reconstructed_tool_call['name']}") + print(f"Reconstructed Arguments: {arguments}") @pytest.fixture @@ -276,9 +334,10 @@ def qwen_tokenizer() -> TokenizerLike: return get_tokenizer("Qwen/Qwen3-32B") -@pytest.fixture -def hermes_parser(qwen_tokenizer: TokenizerLike) -> Hermes2ProToolParser: - return Hermes2ProToolParser(qwen_tokenizer) +@pytest.fixture(params=CONFIGS.keys()) +def hermes_parser(request, qwen_tokenizer: TokenizerLike) -> ToolParser: + config = CONFIGS[request.param] + return config["tool_parser"](qwen_tokenizer) @pytest.fixture @@ -292,7 +351,7 @@ def any_chat_request() -> ChatCompletionRequest: def test_hermes_parser_streaming_just_forward_text( qwen_tokenizer: TokenizerLike, - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: text = """This is some prior text that has nothing to do with tool calling.""" @@ -324,7 +383,7 @@ def test_hermes_parser_streaming_just_forward_text( def test_hermes_parser_streaming_failure_case_bug_19056( qwen_tokenizer: TokenizerLike, - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: text = """ @@ -358,7 +417,7 @@ def test_hermes_parser_streaming_failure_case_bug_19056( def test_hermes_parser_streaming( qwen_tokenizer: TokenizerLike, - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: text = '\ @@ -387,16 +446,20 @@ def test_hermes_parser_streaming( delta_messages.append(delta) print(delta_messages) assert delta_messages[0].tool_calls[0].function.name == "get_current_temperature" - tool_call_args = "".join( - delta.tool_calls[0].function.arguments or "" for delta in delta_messages - ) - assert tool_call_args == ( - '{"location":"San Francisco, California, United States", "unit": "celsius"}' + # load to normalize whitespace + tool_call_args = json.loads( + "".join( + delta.tool_calls[0].function.arguments or "" for delta in delta_messages + ) ) + assert tool_call_args == { + "location": "San Francisco, California, United States", + "unit": "celsius", + } def test_hermes_parser_non_streaming_no_tool_call( - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: text = """This is not a tool call.""" @@ -410,7 +473,7 @@ def test_hermes_parser_non_streaming_no_tool_call( def test_hermes_parser_non_streaming_tool_call_between_tags( - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: text = """ @@ -428,9 +491,12 @@ def test_hermes_parser_non_streaming_tool_call_between_tags( def test_hermes_parser_non_streaming_tool_call_until_eos( - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: + if isinstance(hermes_parser, Granite4ToolParser): + pytest.skip(reason="The Granite4 tool parser enforces a complete response") + text = """ {"name": "final_answer", "arguments": {"trigger": true}}""" tool_call = hermes_parser.extract_tool_calls( @@ -445,7 +511,7 @@ def test_hermes_parser_non_streaming_tool_call_until_eos( def test_hermes_parser_non_streaming_tool_call_invalid_json( - hermes_parser: Hermes2ProToolParser, + hermes_parser: ToolParser, any_chat_request: ChatCompletionRequest, ) -> None: # Missing closing brace to trigger exception diff --git a/tests/kernels/core/test_fused_quant_layernorm.py b/tests/kernels/core/test_fused_quant_layernorm.py index fe06605af25d..f9c01f4f1e62 100644 --- a/tests/kernels/core/test_fused_quant_layernorm.py +++ b/tests/kernels/core/test_fused_quant_layernorm.py @@ -280,21 +280,22 @@ def test_rms_norm( assert torch.allclose(ref_residual, ops_residual) output = torch.empty(x.shape, dtype=quant_dtype, device=x.device) - scales = torch.empty( - (x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32 - ) - if group_size is None: + scales = torch.empty( + (x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32 + ) opcheck( torch.ops._C.rms_norm_dynamic_per_token_quant, (output, x, layer.weight, scales, 1e-5, scale_ub, residual), ) else: - # TODO(luka/eliza) opcheck is broken? - # Somehow the cloned args are getting mutated in-place, - # which causes the opcheck to fail. - # https://github.com/vllm-project/vllm/issues/36688 - return + assert hidden_size % group_size[1] == 0 + num_groups = hidden_size // group_size[1] + scales = torch.empty( + (num_groups, num_tokens), + device=x.device, + dtype=torch.float32, + ).transpose(0, 1) opcheck( torch.ops._C.rms_norm_per_block_quant, ( diff --git a/tests/kernels/quantization/test_nvfp4_quant.py b/tests/kernels/quantization/test_nvfp4_quant.py index 1d2f9d413044..e2db5975882e 100644 --- a/tests/kernels/quantization/test_nvfp4_quant.py +++ b/tests/kernels/quantization/test_nvfp4_quant.py @@ -159,6 +159,52 @@ def test_quantize_to_fp4( torch.testing.assert_close(scale_ans, scale_ref) +@pytest.mark.parametrize( + "shape", + [(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)], +) +@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False]) +@torch.inference_mode() +def test_python_util_matches_cpp_allocation( + shape: tuple[int, int], + is_sf_swizzled_layout: bool, +) -> None: + """ + Verify that the Python utility (create_fp4_output_tensors) allocates + tensors with the same shapes and dtypes as the C++ functional variant + (scaled_fp4_quant_func). + """ + from vllm._custom_ops import create_fp4_output_tensors + + torch.set_default_device("cuda:0") + m, n = shape + input_tensor = torch.randn((m, n), dtype=torch.bfloat16) + input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0") + + # C++ functional variant allocates internally + cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant( + input_tensor, input_scale, is_sf_swizzled_layout + ) + + # Python utility + py_out, py_scale = create_fp4_output_tensors( + m, n, torch.device("cuda:0"), is_sf_swizzled_layout + ) + + assert py_out.shape == cpp_out.shape, ( + f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}" + ) + assert py_out.dtype == cpp_out.dtype, ( + f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}" + ) + assert py_scale.shape == cpp_scale.shape, ( + f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}" + ) + assert py_scale.dtype == cpp_scale.dtype, ( + f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}" + ) + + @pytest.mark.parametrize("pad_shape", PAD_SHAPES) @torch.inference_mode() def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None: diff --git a/tests/models/multimodal/pooling/test_clip.py b/tests/models/multimodal/pooling/test_clip.py index 95c678558f4f..14ede6c1d328 100644 --- a/tests/models/multimodal/pooling/test_clip.py +++ b/tests/models/multimodal/pooling/test_clip.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch from transformers import CLIPModel from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner @@ -50,13 +51,16 @@ def _run_test( if "pixel_values" in inputs: pooled_output = hf_model.model.get_image_features( pixel_values=inputs.pixel_values, - ).squeeze(0) + ) else: pooled_output = hf_model.model.get_text_features( input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, - ).squeeze(0) + ) + if not isinstance(pooled_output, torch.Tensor): + pooled_output = pooled_output.pooler_output + pooled_output = pooled_output.squeeze(0) all_outputs.append(pooled_output.tolist()) hf_outputs = all_outputs diff --git a/tests/models/multimodal/pooling/test_siglip.py b/tests/models/multimodal/pooling/test_siglip.py index 0b8cd33ccfb9..4617250e38f4 100644 --- a/tests/models/multimodal/pooling/test_siglip.py +++ b/tests/models/multimodal/pooling/test_siglip.py @@ -4,6 +4,7 @@ from typing import Any import pytest +import torch from transformers import SiglipModel from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner @@ -68,12 +69,15 @@ def _run_test( if "pixel_values" in inputs: pooled_output = hf_model.model.get_image_features( pixel_values=inputs.pixel_values, - ).squeeze(0) + ) else: pooled_output = hf_model.model.get_text_features( input_ids=inputs.input_ids, - ).squeeze(0) + ) + if not isinstance(pooled_output, torch.Tensor): + pooled_output = pooled_output.pooler_output + pooled_output = pooled_output.squeeze(0) all_outputs.append(pooled_output.tolist()) hf_outputs = all_outputs diff --git a/tests/models/quantization/test_mxfp8.py b/tests/models/quantization/test_mxfp8.py new file mode 100644 index 000000000000..2cb0f2008878 --- /dev/null +++ b/tests/models/quantization/test_mxfp8.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E tests for online MXFP8 quantization. + +Loads a BF16 model with ``--quantization mxfp8`` (online quantization) and +compares log-probabilities against the same model served in BF16 without +quantization. This exercises the full pipeline: config parsing, +``Mxfp8OnlineLinearMethod``, ``Mxfp8OnlineMoEMethod``, weight loading, +online quantization / shuffling, and inference through ``apply_monolithic``. + +Layer skipping (``modules_to_not_convert``) is configured in the model's +``config.json`` under ``quantization_config`` and is not tested here. + +``example_prompts`` is a pytest fixture (from conftest.py) that loads 8 +diverse prompts from ``tests/prompts/example.txt``. +""" + +import pytest + +from tests.quantization.utils import is_quant_method_supported + +from ..utils import check_logprobs_close + +# A small MoE model that fits on a single GPU and has both linear + MoE layers. +MOE_MODEL = "Qwen/Qwen3-30B-A3B" +# A small dense model (no MoE) to validate the linear-only path. +DENSE_MODEL = "Qwen/Qwen3-0.6B" + +MAX_MODEL_LEN = 1024 +MAX_TOKENS = 4 +NUM_LOG_PROBS = 8 + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.quant_model +@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"]) +def test_mxfp8_logprobs( + vllm_runner, + example_prompts, + model: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Compare BF16 baseline logprobs against online MXFP8-quantized model. + + Runs the same model twice -- once in BF16 (baseline) and once with + online MXFP8 quantization -- then checks that the top log-probabilities + are close. Only 4 tokens are generated to keep the test fast while + still catching numerical divergence. + """ + with monkeypatch.context() as m: + m.setenv("TOKENIZERS_PARALLELISM", "true") + + with vllm_runner( + model, + max_model_len=MAX_MODEL_LEN, + enforce_eager=True, + ) as vllm_model: + baseline_outputs = vllm_model.generate_greedy_logprobs( + example_prompts, MAX_TOKENS, NUM_LOG_PROBS + ) + + with vllm_runner( + model, + max_model_len=MAX_MODEL_LEN, + enforce_eager=True, + quantization="mxfp8", + ) as vllm_model: + test_outputs = vllm_model.generate_greedy_logprobs( + example_prompts, MAX_TOKENS, NUM_LOG_PROBS + ) + + check_logprobs_close( + outputs_0_lst=baseline_outputs, + outputs_1_lst=test_outputs, + name_0="bf16", + name_1="mxfp8", + ) + + +@pytest.mark.skipif( + not is_quant_method_supported("mxfp8"), + reason="mxfp8 is not supported on this GPU type (requires sm_100+).", +) +@pytest.mark.quant_model +@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"]) +def test_mxfp8_generation(vllm_runner, model: str) -> None: + """Smoke test: verify online MXFP8 model generates coherent text.""" + prompt = "1 2 3 4 5" + with vllm_runner( + model, + enforce_eager=True, + quantization="mxfp8", + max_model_len=MAX_MODEL_LEN, + ) as vllm_model: + output = vllm_model.generate_greedy([prompt], max_tokens=5) + + generated = output[0][1] + assert len(generated) > len(prompt), ( + f"MXFP8 model produced no new tokens. Output: {generated!r}" + ) diff --git a/tests/reasoning/test_gptoss_reasoning_parser.py b/tests/reasoning/test_gptoss_reasoning_parser.py index 6013fa642edd..3b1327acb688 100644 --- a/tests/reasoning/test_gptoss_reasoning_parser.py +++ b/tests/reasoning/test_gptoss_reasoning_parser.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from unittest.mock import Mock + import pytest from transformers import AutoTokenizer +from vllm.entrypoints.mcp.tool_server import ToolServer from vllm.reasoning import ReasoningParser -from vllm.reasoning.gptoss_reasoning_parser import GptOssReasoningParser +from vllm.reasoning.gptoss_reasoning_parser import ( + GptOssReasoningParser, + from_builtin_tool_to_tag, + no_func_reasoning_tag, +) REASONING_MODEL_NAME = "openai/gpt-oss-120b" @@ -142,3 +150,133 @@ def test_gptoss_is_reasoning_end( output_ids = gpt_oss_tokenizer.convert_tokens_to_ids(output) actual_is_reasoning_end = parser.is_reasoning_end(output_ids) assert is_reasoning_end == actual_is_reasoning_end + + +class TestGptOssStructuralTags: + """Test cases for GptOssReasoningParser structural tag functionality.""" + + @pytest.fixture + def mock_tokenizer(self): + """Create a mock tokenizer for testing.""" + tokenizer = Mock() + tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5]) + tokenizer.get_vocab = Mock(return_value={"<|end|>": 6}) + return tokenizer + + @pytest.fixture + def reasoning_parser(self, mock_tokenizer): + """Create a GptOssReasoningParser instance.""" + return GptOssReasoningParser(mock_tokenizer) + + def test_prepare_structured_tag_no_tool_server(self, reasoning_parser): + """Test prepare_structured_tag with no tool server.""" + result = reasoning_parser.prepare_structured_tag(None, None) + expected = json.dumps(no_func_reasoning_tag) + + assert result == expected + + # Verify the structure is correct + parsed = json.loads(result) + assert parsed["type"] == "structural_tag" + assert parsed["format"]["type"] == "triggered_tags" + assert len(parsed["format"]["tags"]) == 1 + assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>" + assert parsed["format"]["triggers"] == ["<|channel|>analysis"] + + def test_prepare_structured_tag_with_original_tag(self, reasoning_parser): + """Test prepare_structured_tag when original_tag is provided.""" + original_tag = '{"custom": "tag"}' + result = reasoning_parser.prepare_structured_tag(original_tag, None) + + # Should return the original tag unchanged + assert result == original_tag + + def test_from_builtin_tool_to_tag(self): + """Test from_builtin_tool_to_tag function.""" + tags = from_builtin_tool_to_tag("python") + + assert len(tags) == 2 + assert tags[0]["begin"] == "<|channel|>commentary to=python" + assert tags[0]["content"]["type"] == "any_text" + assert tags[0]["end"] == "<|end|>" + + assert tags[1]["begin"] == "<|channel|>analysis to=python" + assert tags[1]["content"]["type"] == "any_text" + assert tags[1]["end"] == "<|end|>" + + @pytest.mark.parametrize( + "tools", + [ + [], + ["browser"], + ["python"], + ["container"], + ["browser", "python"], + ["browser", "container"], + ["python", "container"], + ["browser", "python", "container"], + ], + ) + def test_json_validity_comprehensive(self, reasoning_parser, tools): + """Test JSON validity across all possible tool combinations.""" + tool_server = Mock(spec=ToolServer) + tool_server.has_tool = Mock(side_effect=lambda tool: tool in tools) + + result = reasoning_parser.prepare_structured_tag(None, tool_server) + parsed_result = json.loads(result) + + assert parsed_result["type"] == "structural_tag" + assert "format" in parsed_result + assert "tags" in parsed_result["format"] + assert "triggers" in parsed_result["format"] + + # Tag count should be: 1 (analysis) + 2 * len(tools) + expected_tag_count = 1 + (2 * len(tools)) + assert len(parsed_result["format"]["tags"]) == expected_tag_count + + # Verify triggers are correctly configured + expected_triggers = ["<|channel|>analysis"] + if tools: + expected_triggers.append("<|channel|>commentary to=") + assert set(parsed_result["format"]["triggers"]) == set(expected_triggers) + + def test_no_cross_request_state_pollution(self, reasoning_parser): + """Test that sequential calls with different tool servers produce + independent results, guarding against shared mutable state + (e.g. missing deepcopy in tag_with_builtin_funcs).""" + tool_server_1 = Mock(spec=ToolServer) + tool_server_1.has_tool = Mock(side_effect=lambda tool: tool == "python") + + tool_server_2 = Mock(spec=ToolServer) + tool_server_2.has_tool = Mock(side_effect=lambda tool: tool == "browser") + + result_1 = reasoning_parser.prepare_structured_tag(None, tool_server_1) + result_2 = reasoning_parser.prepare_structured_tag(None, tool_server_2) + + tags_1 = [tag["begin"] for tag in json.loads(result_1)["format"]["tags"]] + tags_2 = [tag["begin"] for tag in json.loads(result_2)["format"]["tags"]] + + assert "<|channel|>commentary to=python" in tags_1 + assert "<|channel|>commentary to=browser" not in tags_1 + + assert "<|channel|>commentary to=browser" in tags_2 + assert "<|channel|>commentary to=python" not in tags_2 + + def test_tag_format_consistency(self, reasoning_parser): + """Test that all generated tags follow consistent format, + catching malformed tags from from_builtin_tool_to_tag.""" + tool_server = Mock(spec=ToolServer) + tool_server.has_tool = Mock( + side_effect=lambda tool: tool in ["python", "browser"] + ) + + result = reasoning_parser.prepare_structured_tag(None, tool_server) + parsed_result = json.loads(result) + + for tag in parsed_result["format"]["tags"]: + assert "begin" in tag + assert "content" in tag + assert "end" in tag + assert tag["content"]["type"] == "any_text" + assert tag["end"] == "<|end|>" + assert tag["begin"].startswith("<|channel|>") diff --git a/tests/v1/distributed/test_internal_lb_dp.py b/tests/v1/distributed/test_internal_lb_dp.py index 8f7459e95ef6..efd9fc607dbb 100644 --- a/tests/v1/distributed/test_internal_lb_dp.py +++ b/tests/v1/distributed/test_internal_lb_dp.py @@ -12,7 +12,7 @@ import pytest_asyncio import requests -from tests.utils import RemoteOpenAIServer +from tests.utils import ROCM_ENV_OVERRIDES, RemoteOpenAIServer from tests.v1.utils import check_request_balancing from vllm.platforms import current_platform @@ -27,6 +27,84 @@ NUM_NODES = 2 +async def _make_completion_request( + client: openai.AsyncOpenAI, + model_name: str, +) -> openai.types.Completion: + """Make a single completion request and validate the response. + + Uses temperature=1.0 to ensure diverse outputs across concurrent + requests for realistic load balancer testing. + """ + completion = await client.completions.create( + model=model_name, + prompt="Hello, my name is", + max_tokens=5, + temperature=1.0, + ) + + assert completion.id is not None, ( + f"Expected non-None completion id. usage={completion.usage!r}" + ) + assert completion.choices is not None and len(completion.choices) == 1, ( + f"Expected 1 choice, got " + f"{len(completion.choices) if completion.choices else 'None'}" + ) + + choice = completion.choices[0] + # With temperature=1.0, the model may emit a stop token immediately, + # producing empty text with finish_reason='stop'. This is valid + # model behavior - the test's purpose is load balancing, not output + # quality. + assert choice.finish_reason in ("length", "stop"), ( + f"Expected finish_reason 'length' or 'stop', " + f"got {choice.finish_reason!r}. text={choice.text!r}" + ) + if choice.finish_reason == "length": + assert len(choice.text) >= 1, ( + f"Expected non-empty text with finish_reason='length', got {choice.text!r}" + ) + + assert completion.usage.prompt_tokens > 0, ( + f"Expected positive prompt_tokens, got {completion.usage.prompt_tokens}" + ) + assert completion.usage.total_tokens > 0, ( + f"Expected positive total_tokens, got {completion.usage.total_tokens}" + ) + return completion + + +async def _run_request_bursts( + client: openai.AsyncOpenAI, + model_name: str, + num_requests: int = 200, + num_bursts: int = 2, +): + """Send multiple bursts of completion requests and validate all succeed.""" + for burst in range(num_bursts): + all_tasks = [] + for _ in range(num_requests): + all_tasks.append( + asyncio.create_task(_make_completion_request(client, model_name)) + ) + await asyncio.sleep(0.01) + + results = await asyncio.gather(*all_tasks, return_exceptions=True) + assert len(results) == num_requests, ( + f"Burst {burst}: expected {num_requests} results, got {len(results)}" + ) + + for result in results: + if isinstance(result, BaseException): + raise result + + assert all(completion is not None for completion in results), ( + f"Burst {burst}: some completions were None" + ) + + await asyncio.sleep(0.5) + + class MultinodeInternalLBServerManager: """Manages multi-node data parallel vLLM server instances for internal load balancer testing using --headless mode.""" @@ -108,6 +186,7 @@ def start_server(sidx: int, r: int, sargs: list[str]): auto_port=False, env_dict={ "VLLM_SERVER_DEV_MODE": "1", + **ROCM_ENV_OVERRIDES, current_platform.device_control_env_var: ",".join( str(current_platform.device_id_to_physical_device_id(i)) for i in range(r, r + gpus_per_node) @@ -229,6 +308,7 @@ def start_api_server(): auto_port=False, env_dict={ "VLLM_SERVER_DEV_MODE": "1", + **ROCM_ENV_OVERRIDES, # No GPUs needed for API-only server }, ) @@ -249,10 +329,11 @@ def start_engines_server(): engines_server_args, auto_port=False, env_dict={ + **ROCM_ENV_OVERRIDES, current_platform.device_control_env_var: ",".join( str(current_platform.device_id_to_physical_device_id(i)) for i in range(self.dp_size * self.tp_size) - ) + ), }, ) server.__enter__() @@ -395,58 +476,15 @@ async def test_multinode_dp_completion( servers: list[tuple[RemoteOpenAIServer, list[str]]], model_name: str, ) -> None: - async def make_request(): - completion = await client.completions.create( - model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0 - ) - - assert completion.id is not None - assert completion.choices is not None and len(completion.choices) == 1 - - choice = completion.choices[0] - # The exact number of tokens can vary slightly with temperature=1.0, - # so we check for a reasonable minimum length. - assert len(choice.text) >= 1 - # Finish reason might not always be 'length' if the model finishes early - # or due to other reasons, especially with high temperature. - # So, we'll accept 'length' or 'stop'. - assert choice.finish_reason in ("length", "stop") - - # Token counts can also vary, so we check they are positive. - assert completion.usage.completion_tokens > 0 - assert completion.usage.prompt_tokens > 0 - assert completion.usage.total_tokens > 0 - return completion - # Test single request - result = await make_request() + result = await _make_completion_request(client, model_name) assert result is not None print("Multi-node internal LB handled single completion request successfully") await asyncio.sleep(0.5) - # Send multiple requests - internal LB should distribute across DP ranks - num_requests = 200 - all_tasks = [] - for _ in range(num_requests): - all_tasks.append(asyncio.create_task(make_request())) - await asyncio.sleep(0.01) - - results = await asyncio.gather(*all_tasks) - assert len(results) == num_requests - assert all(completion is not None for completion in results) - - await asyncio.sleep(0.5) - - # Second burst of requests - all_tasks = [] - for _ in range(num_requests): - all_tasks.append(asyncio.create_task(make_request())) - await asyncio.sleep(0.01) - - results = await asyncio.gather(*all_tasks) - assert len(results) == num_requests - assert all(completion is not None for completion in results) + # Send multiple bursts - internal LB should distribute across DP ranks + await _run_request_bursts(client, model_name) _, server_args = servers[0] api_server_count = ( @@ -570,59 +608,16 @@ async def test_api_only_multinode_dp_completion( ) -> None: """Test API-only server with all engines on separate headless server.""" - async def make_request(): - completion = await api_only_client.completions.create( - model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0 - ) - - assert completion.id is not None - assert completion.choices is not None and len(completion.choices) == 1 - - choice = completion.choices[0] - # The exact number of tokens can vary slightly with temperature=1.0, - # so we check for a reasonable minimum length. - assert len(choice.text) >= 1 - # Finish reason might not always be 'length' if the model finishes - # early or due to other reasons, especially with high temperature. - # So, we'll accept 'length' or 'stop'. - assert choice.finish_reason in ("length", "stop") - - # Token counts can also vary, so we check they are positive. - assert completion.usage.completion_tokens > 0 - assert completion.usage.prompt_tokens > 0 - assert completion.usage.total_tokens > 0 - return completion - # Test single request - result = await make_request() + result = await _make_completion_request(api_only_client, model_name) assert result is not None print("API-only server handled single completion request successfully") await asyncio.sleep(0.5) - # Send multiple requests - should be distributed across engines on + # Send multiple bursts - should be distributed across engines on # headless server - num_requests = 200 - all_tasks = [] - for _ in range(num_requests): - all_tasks.append(asyncio.create_task(make_request())) - await asyncio.sleep(0.01) - - results = await asyncio.gather(*all_tasks) - assert len(results) == num_requests - assert all(completion is not None for completion in results) - - await asyncio.sleep(0.5) - - # Second burst of requests - all_tasks = [] - for _ in range(num_requests): - all_tasks.append(asyncio.create_task(make_request())) - await asyncio.sleep(0.01) - - results = await asyncio.gather(*all_tasks) - assert len(results) == num_requests - assert all(completion is not None for completion in results) + await _run_request_bursts(api_only_client, model_name) api_server, api_server_args = api_only_servers[0] api_server_count = ( diff --git a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh index 684e2ec4d7b9..245b5473448a 100755 --- a/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh @@ -18,11 +18,19 @@ dp_ep_configs=( "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1) "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1) ) +hybrid_ssm_configs=( + "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code" + # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models. + "ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling" +) # Select config array based on DP_EP env var if [[ -n "${DP_EP:-}" ]]; then configs=("${dp_ep_configs[@]}") echo "DP_EP is set, using dp_ep_configs" +elif [[ -n "${HYBRID_SSM:-}" ]]; then + configs=("${hybrid_ssm_configs[@]}") + echo "HYBRID_SSM is set, using hybrid_ssm_configs." else configs=("${tp_configs[@]}") fi diff --git a/tests/v1/kv_connector/nixl_integration/test_accuracy.py b/tests/v1/kv_connector/nixl_integration/test_accuracy.py index 674e65c25ef4..a7fea4e630c9 100644 --- a/tests/v1/kv_connector/nixl_integration/test_accuracy.py +++ b/tests/v1/kv_connector/nixl_integration/test_accuracy.py @@ -18,6 +18,7 @@ "deepseek-ai/deepseek-vl2-tiny": 0.19, "deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65, "google/gemma-3-4b-it": 0.74, + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84, } SIMPLE_PROMPT = ( diff --git a/tests/v1/kv_connector/unit/test_kv_cache_layout.py b/tests/v1/kv_connector/unit/test_kv_cache_layout.py new file mode 100644 index 000000000000..7f8028991703 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_kv_cache_layout.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +def test_mla_backend_rejects_cross_layer_kv_cache(): + """MLA backends return identity permutation (layers dim first) + to signal cross-layer KV cache is unsupported.""" + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonBackend, + ) + + stride_order = MLACommonBackend.get_kv_cache_stride_order( + include_num_layers_dimension=True + ) + assert stride_order == (0, 1, 2, 3) + assert stride_order[0] == 0 # layers dim first => no cross-layer + assert MLACommonBackend.get_kv_cache_stride_order( + include_num_layers_dimension=False + ) == (0, 1, 2) + + +def test_deepseek_v32_indexer_rejects_cross_layer_kv_cache(): + """DeepseekV32Indexer returns identity permutation (layers dim first) + to signal cross-layer KV cache is unsupported.""" + from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerBackend, + ) + + stride_order = DeepseekV32IndexerBackend.get_kv_cache_stride_order( + include_num_layers_dimension=True + ) + assert stride_order == (0, 1, 2, 3) + assert stride_order[0] == 0 # layers dim first => no cross-layer + assert DeepseekV32IndexerBackend.get_kv_cache_stride_order( + include_num_layers_dimension=False + ) == (0, 1, 2) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 2ee224013131..902957e18309 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -84,10 +84,13 @@ def mock_parallel_groups(): yield mock_group -def _setup_kv_transfer_request(request, remote_host="127.0.0.1", fake_port=4789): +def _setup_kv_transfer_request( + request, remote_host="127.0.0.1", fake_port=4789, fake_transfer_id="0" +): """Setup KV transfer parameters for a request.""" request.kv_transfer_params.update( { + "transfer_id": fake_transfer_id, "remote_notify_port": fake_port, "remote_block_ids": None, "remote_host": remote_host, diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 5dd90eb50f1a..095bd4c3dd98 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -53,7 +53,13 @@ from vllm.v1.attention.backends.utils import set_kv_cache_layout from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.output_processor import OutputProcessor -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, +) from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import RequestStatus from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin @@ -332,8 +338,20 @@ def test_kv_transfer_handshake(dist_init): # Prefill connector will register KV cache to populate proper handshake # metadata. - # TODO this must match with values used in kv cache config - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2) + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0", "layer1", "layer2"], + FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, + ), + ) + ] + kv_cache_config = KVCacheConfig( + num_blocks=2, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups + ) prefill_connector = NixlConnector( vllm_config, KVConnectorRole.WORKER, kv_cache_config ) @@ -437,7 +455,7 @@ def __init__( self.kv_cache_layout = kv_cache_layout # Mock register_kv_caches attribute needed for tests that do not call it. self.src_xfer_handles_by_block_size = {self.block_size: 1} - test_shape = self.attn_backend.get_kv_cache_shape( + test_shape = self.attn_backends[0].get_kv_cache_shape( num_blocks=1, block_size=16, num_kv_heads=1, head_size=1 ) self.kv_topo = TpKVTopology( @@ -447,7 +465,7 @@ def __init__( remote_block_size=self._block_size, # shared state is_mla=self.use_mla, total_num_kv_heads=self.model_config.get_total_num_kv_heads(), - attn_backend=self.attn_backend, + attn_backends=self.attn_backends, tensor_shape=test_shape, ) @@ -501,6 +519,7 @@ def _nixl_handshake( # is started. We mock HND here. kv_cache_layout="HND", block_size=self.block_size, + ssm_sizes=(0, 0), ), remote_tp_rank=remote_tp_rank, remote_tp_size=remote_tp_size, @@ -951,6 +970,7 @@ def test_handshake_fails_on_kv_cache_layout_mismatch( block_lens=worker.block_len_per_layer, kv_cache_layout=mismatched_layout, block_size=worker.block_size, + ssm_sizes=(0, 0), ) with pytest.raises(RuntimeError): @@ -1006,6 +1026,7 @@ def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental( block_lens=[i * 2 for i in worker.block_len_per_layer], kv_cache_layout="HND", block_size=worker.block_size, + ssm_sizes=(0, 0), ) # We don't check layout for homogeneous TP and MLA for now, as the @@ -1496,9 +1517,47 @@ def test_register_kv_caches( # test run if not mocking. mock_get_attn_backend.return_value = backend_cls mock_get_attn_backends.return_value = [backend_cls] - + num_layers = 32 + block_size = 16 + num_blocks = 8 + num_heads = 4 + head_size = 16 + + # TODO (NickLucche) the fact that connector depends on kv_cache_config for init + # but cross-layer preference cant be inferred prior to creating kv_cache_config + # is a bit awkward. + dummy_connector = NixlConnector( + vllm_config, + KVConnectorRole.WORKER, + make_kv_cache_config(block_size=block_size), + ) + kv_cache_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_heads, + head_size=head_size, + dtype=torch.float16, + ) + if dummy_connector.prefer_cross_layer_blocks: + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=kv_cache_spec.page_size_bytes * num_blocks, + shared_by=["all-layers"], + ) + for _ in range(num_layers) + ], + kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)], + ) + else: + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec) + ], + ) # Create connector - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2) connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config) connector.connector_worker = FakeNixlConnectorWorker( vllm_config, @@ -1526,35 +1585,6 @@ def test_register_kv_caches( or connector.prefer_cross_layer_blocks ) if connector.prefer_cross_layer_blocks: - num_layers = 32 - block_size = 16 - num_blocks = 8 - # Keep the fake worker's expected num_blocks in sync with the - # cross-layer tensor we are about to register. - worker_kv_cache_config = make_kv_cache_config( - block_size=block_size, num_blocks=num_blocks - ) - connector.connector_worker.kv_cache_config = worker_kv_cache_config - connector.connector_worker.num_blocks = worker_kv_cache_config.num_blocks - kv_cache_spec = AttentionSpec( - block_size=block_size, - num_kv_heads=4, - head_size=64, - dtype=torch.bfloat16, - ) - kv_cache_config = KVCacheConfig( - num_blocks=num_blocks, - kv_cache_tensors=[ - KVCacheTensor( - size=kv_cache_spec.page_size_bytes * num_blocks, - shared_by=["dummy-layer"], - ) - for i in range(num_layers) - ], - # allocate_uniform_kv_caches does not use this - kv_cache_groups=[], - ) - with set_current_vllm_config(vllm_config): _, cross_layers_kv_cache, _ = ( KVConnectorModelRunnerMixin.allocate_uniform_kv_caches( @@ -1586,12 +1616,8 @@ def test_register_kv_caches( expected_blocks_count = 8 kv_caches = {"all-layers": cross_layers_kv_cache} - else: # Create test kv cache tensors using proper backend shape - kv_cache_spec = cast( - AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec - ) kv_cache_shape = backend_cls.get_kv_cache_shape( num_blocks=kv_cache_config.num_blocks, block_size=kv_cache_spec.block_size, @@ -2261,7 +2287,7 @@ def test_compatibility_hash_validation( kv_cache_spec = cast( AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec ) - kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape( + kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape( num_blocks=kv_cache_config.num_blocks, block_size=kv_cache_spec.block_size, num_kv_heads=kv_cache_spec.num_kv_heads, @@ -2269,10 +2295,14 @@ def test_compatibility_hash_validation( ) shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype) + # Build kv_caches from the actual layer names in kv_cache_config so that + # _layer_specs lookups in register_kv_caches always find a matching key. + layer_names = [ + name for group in kv_cache_config.kv_cache_groups for name in group.layer_names + ] kv_caches = { - "layer0": shared_tensor, - "layer1": unique_tensor, - "layer2": shared_tensor, + name: shared_tensor if i % 2 == 0 else unique_tensor + for i, name in enumerate(layer_names) } decode_connector.register_kv_caches(kv_caches) @@ -2312,6 +2342,7 @@ def test_compatibility_hash_validation( block_lens=[4096 * prefill_block_size], # slot_size * block_size kv_cache_layout="HND", block_size=prefill_block_size, + ssm_sizes=(0, 0), ) handshake_payload = NixlHandshakePayload( compatibility_hash=remote_hash, @@ -2391,7 +2422,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario) remote_block_size=decode_worker._block_size, # shared state is_mla=decode_worker.use_mla, total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(), - attn_backend=backend, + attn_backends=[backend], tensor_shape=test_shape, ) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 636d51402bde..d4b0c28a5de5 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -74,6 +74,8 @@ def test_logical_to_kernel_block_ids_with_hma(): # Simulate HMA scenario: logical block size = 32, kernel block size = 16 # So each logical block maps to 2 kernel blocks eg [0]->[0,1] worker._physical_blocks_per_logical_kv_block = 2 + # FA + SW groups (neither is MambaSpec, so both get expanded) + worker.kv_cache_config = make_kv_cache_config(block_size=16, hma_enabled=True) # Test conversion: FA + SW group logical_block_ids = [[0, 1, 2], [3, 4]] @@ -201,3 +203,113 @@ def test_nixl_metadata_hma_block_ids_structure(): assert len(req_meta.remote.block_ids) == 2 assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17] assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21] + + +@pytest.mark.cpu_test +def test_get_block_descs_ids_hybrid_ssm(): + """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM + when ratio=1 (no kernel block size mismatch).""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + NixlConnectorWorker, + ) + + worker = object.__new__(NixlConnectorWorker) + + num_blocks = 100 + engine_id = "test-engine" + worker.num_regions = 2 + worker.dst_num_blocks = {engine_id: num_blocks} + worker._has_mamba = True + worker._is_mamba_group = [False, True] + worker._physical_blocks_per_logical_kv_block = 1 + # num_descs = num_regions * num_blocks (no blocks_first doubling) + worker.num_descs = 2 * num_blocks + + fa_blocks = [3, 5] + ssm_blocks = [1, 2] + result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks)) + + # FA group: stride=num_blocks=100, offset=0 + # region0: [3, 5], region1: [103, 105] + # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1), + # offset=num_descs=200 + # region0: [201, 202], region1: [301, 302] + expected = [3, 5, 103, 105, 201, 202, 301, 302] + assert list(result) == expected, f"Expected {expected}, got {list(result)}" + + +@pytest.mark.cpu_test +def test_get_block_descs_ids_kernel_block_mismatch(): + """Test _get_block_descs_ids uses different strides for FA (kernel blocks) + vs SSM (logical blocks) when ratio > 1.""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + NixlConnectorWorker, + ) + + worker = object.__new__(NixlConnectorWorker) + + ratio = 4 + logical_blocks = 100 + num_blocks = logical_blocks * ratio # 400 kernel blocks + engine_id = "test-engine" + worker.num_regions = 2 + worker.dst_num_blocks = {engine_id: num_blocks} + worker._has_mamba = True + worker._is_mamba_group = [False, True] + worker._physical_blocks_per_logical_kv_block = ratio + worker.num_descs = 2 * num_blocks # 800 + + fa_blocks = [3, 7] # kernel-level block IDs + ssm_blocks = [1, 2] # logical block IDs + result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks)) + + # FA group: stride=num_blocks=400, offset=0 + # region0: [3, 7], region1: [403, 407] + # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800 + # region0: [801, 802], region1: [901, 902] + expected = [3, 7, 403, 407, 801, 802, 901, 902] + assert list(result) == expected, f"Expected {expected}, got {list(result)}" + + +@pytest.mark.cpu_test +def test_nixl_metadata_hybrid_ssm_block_ids(): + """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM + groups with different block counts (kernel mismatch active).""" + from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import ( + NixlConnectorMetadata, + ) + + metadata = NixlConnectorMetadata() + + # FA: 8 kernel blocks (2 logical * ratio=4), SSM: 2 logical blocks + fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7] + ssm_blocks = [0, 1] + + metadata.add_new_req_to_recv( + request_id="test-req-hybrid", + local_block_ids=(fa_blocks, ssm_blocks), + kv_transfer_params={ + "remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [20, 21]), + "remote_engine_id": "remote-engine", + "remote_request_id": "prefill-test-req-hybrid", + "remote_host": "localhost", + "remote_port": 1234, + "tp_size": 1, + }, + ) + + assert "test-req-hybrid" in metadata.reqs_to_recv + req_meta = metadata.reqs_to_recv["test-req-hybrid"] + + # Verify local block IDs: different lengths per group + assert len(req_meta.local_block_ids) == 2 + assert list(req_meta.local_block_ids[0]) == fa_blocks + assert list(req_meta.local_block_ids[1]) == ssm_blocks + assert len(req_meta.local_block_ids[0]) != len(req_meta.local_block_ids[1]) + + # Verify remote block IDs: same asymmetry preserved + assert req_meta.remote is not None + assert len(req_meta.remote.block_ids) == 2 + assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17] + assert list(req_meta.remote.block_ids[1]) == [20, 21] + assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1]) diff --git a/tests/v1/structured_output/test_gptoss_structural_tags.py b/tests/v1/structured_output/test_gptoss_structural_tags.py deleted file mode 100644 index fb1eae53d16d..000000000000 --- a/tests/v1/structured_output/test_gptoss_structural_tags.py +++ /dev/null @@ -1,173 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Unit tests for GPT-OSS structural tag support in reasoning (PR #25515).""" - -import json -from unittest.mock import Mock - -import pytest - -from vllm.entrypoints.mcp.tool_server import ToolServer -from vllm.reasoning.gptoss_reasoning_parser import ( - GptOssReasoningParser, - from_builtin_tool_to_tag, - no_func_reaonsing_tag, - tag_with_builtin_funcs, -) - - -class TestGptOssReasoningParser: - """Test cases for GptOssReasoningParser structural tag functionality.""" - - @pytest.fixture - def mock_tokenizer(self): - """Create a mock tokenizer for testing.""" - tokenizer = Mock() - tokenizer.encode = Mock(return_value=[1, 2, 3, 4, 5]) - tokenizer.get_vocab = Mock(return_value={"<|end|>": 6}) - return tokenizer - - @pytest.fixture - def reasoning_parser(self, mock_tokenizer): - """Create a GptOssReasoningParser instance.""" - return GptOssReasoningParser(mock_tokenizer) - - @pytest.fixture - def mock_tool_server_empty(self): - """Create a mock ToolServer with no tools.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(return_value=False) - return tool_server - - @pytest.fixture - def mock_tool_server_with_browser(self): - """Create a mock ToolServer with browser tool.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=lambda tool: tool == "browser") - return tool_server - - @pytest.fixture - def mock_tool_server_with_all_tools(self): - """Create a mock ToolServer with all builtin tools.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock( - side_effect=lambda tool: tool in ["browser", "python", "container"] - ) - return tool_server - - def test_prepare_structured_tag_no_tool_server(self, reasoning_parser): - """Test prepare_structured_tag with no tool server.""" - result = reasoning_parser.prepare_structured_tag(None, None) - expected = json.dumps(no_func_reaonsing_tag) - - assert result == expected - - # Verify the structure is correct - parsed = json.loads(result) - assert parsed["type"] == "structural_tag" - assert parsed["format"]["type"] == "triggered_tags" - assert len(parsed["format"]["tags"]) == 1 - assert parsed["format"]["tags"][0]["begin"] == "<|channel|>analysis<|message|>" - assert parsed["format"]["triggers"] == ["<|channel|>analysis"] - - def test_prepare_structured_tag_with_all_tools( - self, reasoning_parser, mock_tool_server_with_all_tools - ): - """Test prepare_structured_tag with all builtin tools.""" - result = reasoning_parser.prepare_structured_tag( - None, mock_tool_server_with_all_tools - ) - parsed = json.loads(result) - - # Should have analysis tag + tags for all 3 tools (2 tags each) - assert len(parsed["format"]["tags"]) == 7 # 1 analysis + 6 tool tags - - # Check all tool tags are present - tag_begins = [tag["begin"] for tag in parsed["format"]["tags"]] - for tool in ["browser", "python", "container"]: - assert f"<|channel|>commentary to={tool}" in tag_begins - assert f"<|channel|>analysis to={tool}" in tag_begins - - def test_prepare_structured_tag_with_original_tag(self, reasoning_parser): - """Test prepare_structured_tag when original_tag is provided.""" - original_tag = '{"custom": "tag"}' - result = reasoning_parser.prepare_structured_tag(original_tag, None) - - # Should return the original tag unchanged - assert result == original_tag - - def test_from_builtin_tool_to_tag(self): - """Test from_builtin_tool_to_tag function.""" - tags = from_builtin_tool_to_tag("python") - - assert len(tags) == 2 - assert tags[0]["begin"] == "<|channel|>commentary to=python" - assert tags[0]["content"]["type"] == "any_text" - assert tags[0]["end"] == "<|end|>" - - assert tags[1]["begin"] == "<|channel|>analysis to=python" - assert tags[1]["content"]["type"] == "any_text" - assert tags[1]["end"] == "<|end|>" - - def test_tag_with_builtin_funcs(self): - """Test tag_with_builtin_funcs function.""" - builtin_tools = ["browser", "python"] - result = tag_with_builtin_funcs(no_func_reaonsing_tag, builtin_tools) - - assert result["type"] == "structural_tag" - # Should have original analysis tag + 2 tags per tool - assert len(result["format"]["tags"]) == 5 # 1 + 2*2 - - # Should have added commentary trigger - assert "<|channel|>commentary to=" in result["format"]["triggers"] - assert "<|channel|>analysis" in result["format"]["triggers"] - - def test_tag_structure_invariants(self): - """Test that the basic tag structure follows expected format.""" - # Test the base no_func_reaonsing_tag structure - assert no_func_reaonsing_tag["type"] == "structural_tag" - assert no_func_reaonsing_tag["format"]["type"] == "triggered_tags" - assert no_func_reaonsing_tag["format"]["stop_after_first"] is False - - # Verify analysis tag structure - analysis_tag = no_func_reaonsing_tag["format"]["tags"][0] - assert analysis_tag["begin"] == "<|channel|>analysis<|message|>" - assert analysis_tag["content"]["type"] == "any_text" - assert analysis_tag["end"] == "<|end|>" - - def test_json_serialization_valid( - self, reasoning_parser, mock_tool_server_with_all_tools - ): - """Test that all generated tags produce valid JSON.""" - # Test with no tool server - result1 = reasoning_parser.prepare_structured_tag(None, None) - json.loads(result1) # Should not raise - - # Test with empty tool server - empty_server = Mock(spec=ToolServer) - empty_server.has_tool = Mock(return_value=False) - result2 = reasoning_parser.prepare_structured_tag(None, empty_server) - json.loads(result2) # Should not raise - - # Test with tools - result3 = reasoning_parser.prepare_structured_tag( - None, mock_tool_server_with_all_tools - ) - json.loads(result3) # Should not raise - - @pytest.mark.parametrize("tool_name", ["browser", "python", "container"]) - def test_single_tool_integration(self, reasoning_parser, tool_name): - """Test integration with individual tools.""" - tool_server = Mock(spec=ToolServer) - tool_server.has_tool = Mock(side_effect=lambda tool: tool == tool_name) - - result = reasoning_parser.prepare_structured_tag(None, tool_server) - parsed = json.loads(result) - - # Should have 1 analysis + 2 tool-specific tags - assert len(parsed["format"]["tags"]) == 3 - - tag_begins = [tag["begin"] for tag in parsed["format"]["tags"]] - assert f"<|channel|>commentary to={tool_name}" in tag_begins - assert f"<|channel|>analysis to={tool_name}" in tag_begins diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index fdc468d3b25d..63f347d89de2 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -29,6 +29,81 @@ def register_fake(fn): from torch.library import impl_abstract as register_fake +# scaled_fp4_quant functional + out variant for torch.compile buffer management + + +def create_fp4_scale_tensor( + m: int, + n: int, + device: torch.device, + is_sf_swizzled_layout: bool, +) -> torch.Tensor: + """ + Allocate the output scale tensor for scaled_fp4_quant. + + When is_sf_swizzled_layout=True, we use rounded values to store the + swizzled scales. Due to the requirement of the Tensor Core, the minimum + tile is 128x4 for the scales. So, we first pad the scales to multiples + of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are + packed into an int32 for every 4 values. More: + https://docs.nvidia.com/cuda/parallel-thread-execution/ + #tcgen05-mma-scale-factor-b-layout-4x + """ + from vllm.utils.math_utils import round_up + + block_size = 16 + if is_sf_swizzled_layout: + rounded_m = round_up(m, 128) + scale_n = n // block_size + rounded_n = round_up(scale_n, 4) + return torch.empty( + (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 + ) + else: + return torch.empty((m, n // block_size), device=device, dtype=torch.uint8) + + +def create_fp4_output_tensors( + m: int, + n: int, + device: torch.device, + is_sf_swizzled_layout: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Allocate both output tensors for scaled_fp4_quant: + (quantized_output, output_scale). + + Must match the C++ scaled_fp4_quant_func allocation exactly. + """ + output = torch.empty((m, n // 2), device=device, dtype=torch.uint8) + output_scale = create_fp4_scale_tensor(m, n, device, is_sf_swizzled_layout) + return output, output_scale + + +if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"): + + @register_fake("_C::scaled_fp4_quant") + def _scaled_fp4_quant_fake( + input: torch.Tensor, + input_scale: torch.Tensor, + is_sf_swizzled_layout: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + n = input.shape[-1] + m = input.numel() // n + return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout) + + @register_fake("_C::scaled_fp4_quant.out") + def _scaled_fp4_quant_out_fake( + input: torch.Tensor, + input_scale: torch.Tensor, + is_sf_swizzled_layout: bool, + *, + output: torch.Tensor, + output_scale: torch.Tensor, + ) -> None: + return None + + # page attention ops def paged_attention_v1( out: torch.Tensor, @@ -1644,7 +1719,6 @@ def scaled_fp4_quant( input = input.reshape(other_dims, input.shape[-1]) m, n = input.shape block_size = 16 - device = input.device assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}." assert input.dtype in (torch.float16, torch.bfloat16), ( @@ -1658,26 +1732,16 @@ def scaled_fp4_quant( input, input_global_scale ) else: - # Two fp4 values will be packed into an uint8. - output = torch.empty((m, n // 2), device=device, dtype=torch.uint8) - if is_sf_swizzled_layout: - # We use the rounded values to store the swizzled values. Due to the - # requirement of the Tensor Core, the minimum tile is 128x4 for the scales. - # So, we first pad the scales to multiples of 128 and 4. Then, the scales - # (in float8_e4m3fn) are packed into an int32 for every 4 values. More: - # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x - round_up = lambda x, y: (x + y - 1) // y * y - rounded_m = round_up(m, 128) - scale_n = n // block_size - rounded_n = round_up(scale_n, 4) - output_scale = torch.empty( - (rounded_m, rounded_n // 4), device=device, dtype=torch.int32 - ) - else: - output_scale = torch.empty((m, n // 16), device=device, dtype=torch.uint8) - - torch.ops._C.scaled_fp4_quant( - output, input, output_scale, input_global_scale, is_sf_swizzled_layout + # Pre-allocate and call .out variant (same behavior as old in-place API) + output, output_scale = create_fp4_output_tensors( + m, n, input.device, is_sf_swizzled_layout + ) + torch.ops._C.scaled_fp4_quant.out( + input, + input_global_scale, + is_sf_swizzled_layout, + output=output, + output_scale=output_scale, ) output_scale = output_scale.view(torch.float8_e4m3fn) diff --git a/vllm/compilation/caching.py b/vllm/compilation/caching.py index 00fb959211fa..2b667344ff37 100644 --- a/vllm/compilation/caching.py +++ b/vllm/compilation/caching.py @@ -307,13 +307,6 @@ def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction num_submods = len(submod_names) num_artifacts = standalone_compile_artifacts.num_artifacts() - logger.info( - "reconstructing serializable fn from standalone compile " - "artifacts. num_artifacts=%d num_submods=%d", - num_artifacts, - num_submods, - ) - with functorch_ctx: fn = reconstruct_serializable_fn_from_mega_artifact( state=state, @@ -324,7 +317,10 @@ def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction ) logger.info( - "reconstructed serializable fn from standalone compile artifacts" + "reconstructed serializable fn from standalone compile " + "artifacts. num_artifacts=%d num_submods=%d", + num_artifacts, + num_submods, ) return fn diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 13e88448c0f1..78841866f752 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -16,7 +16,11 @@ from vllm.compilation.monitor import validate_cudagraph_capturing_enabled from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id -from vllm.forward_context import BatchDescriptor, get_forward_context +from vllm.forward_context import ( + BatchDescriptor, + get_forward_context, + is_forward_context_available, +) from vllm.logger import init_logger from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform @@ -224,6 +228,12 @@ def clear_graphs(self) -> None: self.concrete_cudagraph_entries.clear() def __call__(self, *args: Any, **kwargs: Any) -> Any | None: + if not is_forward_context_available(): + # No forward context means we are outside the normal + # inference path (e.g. a vision encoder forward pass). + # Just run the underlying function without cudagraphs. + return self.runnable(*args, **kwargs) + forward_context = get_forward_context() batch_descriptor = forward_context.batch_descriptor cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode diff --git a/vllm/compilation/passes/fusion/act_quant_fusion.py b/vllm/compilation/passes/fusion/act_quant_fusion.py index e141003849ac..911775f69967 100644 --- a/vllm/compilation/passes/fusion/act_quant_fusion.py +++ b/vllm/compilation/passes/fusion/act_quant_fusion.py @@ -148,11 +148,11 @@ def pattern( result_silu_mul = self.silu_and_mul_matcher(input) at = auto_functionalized( self.QUANT_OP, - output=result, input=result_silu_mul, - output_scale=output_scale, input_scale=scale, is_sf_swizzled_layout=True, + output=result, + output_scale=output_scale, ) return at[1], at[2] diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index 44dc3d67bb98..f141a7c171f7 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -47,7 +47,7 @@ pass if hasattr(torch.ops._C, "scaled_fp4_quant"): - STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default + STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.out # Max size of the input tensor per world size per device capability # to use flashinfer fused allreduce @@ -562,11 +562,11 @@ def pattern( rms = self.rmsnorm_matcher(all_reduce, weight) quant_out_tuple = auto_functionalized( STATIC_FP4_QUANT_OP, - output=quant_result, input=rms, - output_scale=output_scale, input_scale=input_global_scale, is_sf_swizzled_layout=True, + output=quant_result, + output_scale=output_scale, ) # quant_out, allreduce_output, output_scale @@ -660,11 +660,11 @@ def pattern( rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual) quant_out_tuple = auto_functionalized( STATIC_FP4_QUANT_OP, - output=quant_result, input=rms, - output_scale=output_scale, input_scale=input_global_scale, is_sf_swizzled_layout=True, + output=quant_result, + output_scale=output_scale, ) # quant_out, allreduce_output, output_scale diff --git a/vllm/compilation/passes/fusion/attn_quant_fusion.py b/vllm/compilation/passes/fusion/attn_quant_fusion.py index 5e6bf28c0041..0e1b846af856 100644 --- a/vllm/compilation/passes/fusion/attn_quant_fusion.py +++ b/vllm/compilation/passes/fusion/attn_quant_fusion.py @@ -250,11 +250,11 @@ def pattern( ) at2 = auto_functionalized( self.QUANT_OP, - output=output_quant, input=attn_out_view, - output_scale=output_scale, input_scale=input_scale, is_sf_swizzled_layout=True, + output=output_quant, + output_scale=output_scale, ) output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE) return at2[1], output_scale_view diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 03f680552c58..ec36c12d1776 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -38,7 +38,7 @@ } if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): - QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default # noqa: E501 + QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out # noqa: E501 if current_platform.is_cuda(): QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 diff --git a/vllm/compilation/passes/fusion/rms_quant_fusion.py b/vllm/compilation/passes/fusion/rms_quant_fusion.py index 2d084783d7d7..95ce7b22e0a3 100644 --- a/vllm/compilation/passes/fusion/rms_quant_fusion.py +++ b/vllm/compilation/passes/fusion/rms_quant_fusion.py @@ -63,7 +63,7 @@ def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor: kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501 } if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"): - QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default + QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out if current_platform.is_cuda(): QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501 diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 3796265ffc0a..f4c70cace264 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -13,6 +13,7 @@ CacheDType = Literal[ "auto", + "float16", "bfloat16", "fp8", "fp8_e4m3", diff --git a/vllm/distributed/kv_transfer/kv_connector/utils.py b/vllm/distributed/kv_transfer/kv_connector/utils.py index 155395e84e11..1f889c6c838a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/utils.py +++ b/vllm/distributed/kv_transfer/kv_connector/utils.py @@ -16,10 +16,12 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.platforms import current_platform from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase + from vllm.v1.kv_cache_interface import KVCacheSpec logger = init_logger(__name__) @@ -328,22 +330,26 @@ class TpKVTopology: remote_tp_size: dict[EngineId, int] is_mla: bool total_num_kv_heads: int - attn_backend: type[AttentionBackend] + attn_backends: list[type[AttentionBackend]] engine_id: EngineId remote_block_size: dict[EngineId, int] tensor_shape: torch.Size | None = None + is_mamba: bool = False def __post_init__(self): # Figure out whether the first dimension of the cache is K/V # or num_blocks. This is used to register the memory regions correctly. - _MOCK_BLOCK_SIZE = 16 - kv_cache_shape = self.attn_backend.get_kv_cache_shape( - num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1 - ) - logger.debug("Test kv_cache_shape: %s", kv_cache_shape) + attn_backend = self.attn_backends[0] + if not self.is_mamba: + _MOCK_BLOCK_SIZE = 16 + kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape( + num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1 + ) + logger.debug("Test kv_cache_shape: %s", kv_cache_shape) # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D], # we just mock num_blocks to 1 for the dimension check below. - self._is_kv_layout_blocks_first = ( + # Hybrid SSM models assume a single blocks_first layout + self._is_kv_layout_blocks_first = self.is_mamba or ( len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1 ) @@ -360,7 +366,7 @@ def __post_init__(self): _MOCK_NUM_LAYERS = 80 kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape try: - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order( + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order( include_num_layers_dimension=self._cross_layers_blocks ) except (AttributeError, NotImplementedError): @@ -483,6 +489,30 @@ def get_target_remote_ranks_from_engine_id( remote_tp_size = self.remote_tp_size[remote_engine_id] return self.get_target_remote_ranks(remote_tp_size) + def get_transfer_cache_regions( + self, cache: torch.Tensor, layer_spec: "KVCacheSpec" + ) -> list[torch.Tensor] | torch.Tensor: + """Return the cache tensor(s) to register as NIXL memory regions, + also accounting for hybrid SSM models specificities. + """ + if isinstance(layer_spec, MambaSpec): + # Register the whole kv cache shared tensor, including SSM/Conv. This is + # similar to FI with the difference that SSM/Conv have different sizes + conv, ssm = cache + return [conv] + + # Check may be hacky but it's matching `_update_hybrid_attention_mamba_layout`. + if self.is_mamba and cache.shape[0] == 2: + # When MAMBA is present, all backends are blocks first, so that blocks + # can be shared between attention layers and mamba layers. Runner + # `_update_hybrid_attention_mamba_layout` already adjusted strides + # for FlashAttn-like backends so its num_blocks first. + # Swap [2<>num_blocks] dims to get required layout for hybrid SSM. + cache = cache.transpose(0, 1) + + # Regular case: backends like FA register K/V in separate regions + return cache if self.split_k_and_v else [cache] + def get_current_attn_backends( vllm_config: VllmConfig, layer_names: list[str] | None = None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index d986f686657f..28b997128d46 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -564,7 +564,7 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): remote_block_size=self._block_size, # shared state is_mla=self.use_mla, total_num_kv_heads=self.model_config.get_total_num_kv_heads(), - attn_backend=backend, + attn_backends=[backend], ) self.async_zmq_ctx = zmq.asyncio.Context() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py index d381b527075f..7651bf988284 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py @@ -59,7 +59,12 @@ from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.attention.backends.utils import get_kv_cache_layout from vllm.v1.core.sched.output import SchedulerOutput -from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, +) from vllm.v1.worker.block_table import BlockTable from vllm.v1.worker.utils import select_common_block_size @@ -159,6 +164,7 @@ class NixlAgentMetadata: block_lens: list[int] kv_cache_layout: str block_size: int + ssm_sizes: tuple[int, int] @dataclass @@ -310,6 +316,15 @@ def add_new_req_to_recv( class NixlConnector(KVConnectorBase_V1, SupportsHMA): @property def prefer_cross_layer_blocks(self) -> bool: + if any( + [ + isinstance(group.kv_cache_spec, MambaSpec) + for group in self.kv_cache_config.kv_cache_groups + ] + ): + # Hybrid SSM models do not yet support cross-layer layout + return False + backend = get_current_attn_backend(self._vllm_config) if backend.get_name() not in ( "FLASH_ATTN", @@ -335,12 +350,9 @@ def __init__( kv_cache_config: "KVCacheConfig", ): super().__init__(vllm_config, role, kv_cache_config) - assert vllm_config.kv_transfer_config is not None assert vllm_config.kv_transfer_config.engine_id is not None - for group in kv_cache_config.kv_cache_groups: - if isinstance(group.kv_cache_spec, MambaSpec): - raise ValueError("NixlConnector does not support Mamba models.") + self.kv_cache_config = kv_cache_config self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id self.kv_transfer_config = vllm_config.kv_transfer_config if role == KVConnectorRole.SCHEDULER: @@ -403,6 +415,14 @@ def build_connector_meta( assert self.connector_scheduler is not None return self.connector_scheduler.build_connector_meta(scheduler_output) + def request_finished( + self, + request: "Request", + block_ids: list[int], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, (block_ids,)) + def request_finished_all_groups( self, request: "Request", @@ -434,11 +454,7 @@ def register_cross_layers_kv_cache( self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend] ): assert self.connector_worker is not None - - cross_layer_name = "ALL_LAYERS" - kv_caches = {cross_layer_name: kv_cache} - - self.connector_worker.register_kv_caches(kv_caches) + self.connector_worker.register_cross_layers_kv_caches(kv_cache) def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp): assert self.connector_worker is not None @@ -962,6 +978,40 @@ def __init__( ) ) self.kv_cache_config = kv_cache_config + self._layer_specs = { + layer: group.kv_cache_spec + for group in kv_cache_config.kv_cache_groups + for layer in group.layer_names + } + self.hma_group_size = len(kv_cache_config.kv_cache_tensors) + + # Mamba metadata + self._is_mamba_group = [ + isinstance(group.kv_cache_spec, MambaSpec) + for group in kv_cache_config.kv_cache_groups + ] + mamba_ssm_size = (0, 0) + self._has_mamba = any(self._is_mamba_group) + if self._has_mamba: + assert self._is_hma_required + mamba_spec = next( + spec + for spec in self._layer_specs.values() + if isinstance(spec, MambaSpec) + ) + conv_nbytes, ssm_nbytes = ( + torch.tensor([], dtype=mamba_spec.dtypes[0]).element_size(), # type: ignore[misc] + torch.tensor([], dtype=mamba_spec.dtypes[1]).element_size(), # type: ignore[misc] + ) + conv_shape, ssm_shape = ( + torch.Size(mamba_spec.shapes[0]), + torch.Size(mamba_spec.shapes[1]), + ) + mamba_ssm_size = ( + conv_shape.numel() * conv_nbytes, + ssm_shape.numel() * ssm_nbytes, + ) + self._mamba_ssm_size = mamba_ssm_size # Agent. non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"] @@ -1106,9 +1156,9 @@ def __init__( # Get the attention backend from the first layer # NOTE (NickLucche) models with multiple backends are not supported yet - self.attn_backend = get_current_attn_backend(vllm_config) + self.attn_backends = get_current_attn_backends(vllm_config) + self.backend_name = self.attn_backends[0].get_name() - self.backend_name = self.attn_backend.get_name() self.kv_cache_layout = get_kv_cache_layout() self.host_buffer_kv_cache_layout = self.kv_cache_layout logger.info("Detected attention backend %s", self.backend_name) @@ -1135,6 +1185,8 @@ def __init__( def _sync_block_size_with_kernel(self) -> None: backends = get_current_attn_backends(self.vllm_config) kernel_block_size = select_common_block_size(self.block_size, backends) + # Number of blocks not accounting for kernel block mismatches + self._logical_num_blocks = self.num_blocks if self.block_size != kernel_block_size: logger.info_once( "User-specified logical block size (%s) does not match" @@ -1428,9 +1480,19 @@ def request_ready(f: Future[Any], entry=(req_id, meta)): fut.add_done_callback(request_ready) + def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None: + """Register a cross-layers KV cache tensor with NIXL. + + `use_uniform_kv_cache()` guarantees a single KV cache group whose + layers all share the same `AttentionSpec`, so any layer name from + `_layer_specs` yields the correct per-layer spec for `page_size_bytes`. + """ + first_layer = next(iter(self._layer_specs)) + # Forwarding a real layer name rather than a synthetic key + self.register_kv_caches({first_layer: kv_cache}) + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): """Register the KV Cache data in nixl.""" - self.kv_topo = TpKVTopology( tp_rank=self.tp_rank, engine_id=self.engine_id, @@ -1438,8 +1500,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): remote_block_size=self._block_size, # shared state is_mla=self.use_mla, total_num_kv_heads=self.model_config.get_total_num_kv_heads(), - attn_backend=self.attn_backend, - tensor_shape=next(iter(kv_caches.values())).shape, + attn_backends=self.attn_backends, + # SSM States come in tuples (ssm, conv) + tensor_shape=next(iter(kv_caches.values())).shape + if not self._has_mamba + else None, + is_mamba=self._has_mamba, ) self.compat_hash = compute_nixl_compatibility_hash( self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks @@ -1481,12 +1547,50 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # to better exploit the memory layout (ie num_blocks is the first dim). tensor_size_bytes = None - # Enable different block lengths for different layers when MLA is used. + # Enable different block lengths for different layers *only* when MLA is used. + # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`. self.block_len_per_layer = list[int]() for layer_name, cache_or_caches in xfer_buffers.items(): - cache_list = ( - cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches] + # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to + # that of FI, with block laid out as in `get_backend_aware_kv_block_len`. + # However, physical page_size may differ when kernel requires a specific + # block size. This leads to SSM and FA layers having different num_blocks. + # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this. + layer_spec = self._layer_specs[layer_name] + if isinstance(layer_spec, UniformTypeKVCacheSpecs): + # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs + layer_spec = layer_spec.kv_cache_specs[layer_name] + cache_list = self.kv_topo.get_transfer_cache_regions( + cache_or_caches, layer_spec + ) + # `layer_spec.page_size_bytes` only accounts for logical page_size, that is + # the page_size assuming constant `self._logical_num_blocks`. + physical_page_size = ( + layer_spec.page_size_bytes + if isinstance(layer_spec, MambaSpec) + else layer_spec.page_size_bytes + // self._physical_blocks_per_logical_kv_block ) + # For when registering multiple tensors eg K/V in separate regions. + physical_page_size = physical_page_size // len(cache_list) + if self.kv_topo._cross_layers_blocks: + # When cross-layers blocks are used, multiply by number of layers + physical_page_size = physical_page_size * len( + self.kv_cache_config.kv_cache_tensors + ) + num_blocks = ( + self._logical_num_blocks + if isinstance(layer_spec, MambaSpec) + else self.num_blocks + ) + # `page_size` accounts for physical blocks, st KVCache is always + # [`num_blocks` * `page_size`] + curr_tensor_size_bytes = num_blocks * physical_page_size + if tensor_size_bytes is None: + tensor_size_bytes = curr_tensor_size_bytes + + # TODO (NickLucche) we could eventually unify how we handle FA/FI regions, + # registering a single tensor for both K/V and splitting logically like FI. for cache in cache_list: base_addr = cache.data_ptr() if base_addr in seen_base_addresses: @@ -1494,27 +1598,27 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): # across groups. This results in skipping all tensors but the ones # pointed to by group0. Also, generally we will have more blocks # per tensor but fewer regions. + logger.debug("Skipping %s because it's already seen", layer_name) continue - logger.debug( "Registering layer %s with cache shape: %s", layer_name, cache.shape ) seen_base_addresses.append(base_addr) - curr_tensor_size_bytes = cache.numel() * cache.element_size() - - if tensor_size_bytes is None: - tensor_size_bytes = curr_tensor_size_bytes + # Only record non-Mamba page sizes. + if isinstance(layer_spec, MambaSpec): + self.block_len_per_layer.append( + physical_page_size // self._physical_blocks_per_logical_kv_block + ) + else: + self.block_len_per_layer.append(physical_page_size) - assert cache.shape[0] == self.num_blocks, ( + assert cache.shape[0] == num_blocks, ( "All kv cache tensors must have the same number of blocks" ) - self.block_len_per_layer.append( - curr_tensor_size_bytes // self.num_blocks - ) - if not self.use_mla: - # Different kv cache shape is not supported by HeteroTP + # Different kv cache shape is not supported by HeteroTP. + # This must also hold true for Mamba-like models. assert tensor_size_bytes == curr_tensor_size_bytes, ( "All kv cache tensors must have the same size" ) @@ -1533,6 +1637,21 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses self.num_regions = len(caches_data) + if self.kv_topo.is_kv_layout_blocks_first: + # NOTE (NickLucche) When FlashInfer is used, memory is registered + # with joint KV for each block. This minimizes the overhead in + # registerMem allowing faster descs queries. In order to be able to + # split on kv_heads dim as required by heterogeneous TP, one must + # be able to index K/V separately. Hence we double the number + # of 'virtual' regions here and halve `block_len` below. + # Similarly for Mamba layers, we register SSM+Conv as a single region and + # then duplicate it logically to be able to index SSM/Conv separately. + self.num_regions *= 2 + + # TODO (NickLucche) Adapt to different descs views (engine_id->tp_rank) to + # support heterogeneous TP. + self.num_descs = self.num_regions * self.num_blocks + descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type) logger.debug("Registering descs: %s", caches_data) self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends) @@ -1542,17 +1661,21 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.device_kv_caches = kv_caches self.dst_num_blocks[self.engine_id] = self.num_blocks - if self.kv_topo.is_kv_layout_blocks_first: - # NOTE (NickLucche) When FlashInfer is used, memory is registered - # with joint KV for each block. This minimizes the overhead in - # registerMem allowing faster descs queries. In order to be able to - # split on kv_heads dim as required by heterogeneous TP, one must - # be able to index K/V separately. Hence we double the number - # of 'virtual' regions here and halve `block_len` below. - self.num_regions *= 2 + if self._has_mamba: + logger.info( + "Hybrid SSM registration: num_blocks=%s, " + "logical_num_blocks=%s, ratio=%s, num_regions=%s, " + "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s", + self.num_blocks, + self._logical_num_blocks, + self._physical_blocks_per_logical_kv_block, + self.num_regions, + self.num_descs, + self._mamba_ssm_size, + set(self.block_len_per_layer), + ) # Register local/src descr for NIXL xfer. - self.seen_base_addresses = seen_base_addresses self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = ( self.register_local_xfer_handler(self.block_size) ) @@ -1569,6 +1692,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): if not self.use_host_buffer else self.host_buffer_kv_cache_layout, block_size=self.block_size, + ssm_sizes=self._mamba_ssm_size, ) # Wrap metadata in payload with hash for defensive decoding assert self.compat_hash is not None @@ -1594,40 +1718,65 @@ def register_local_xfer_handler( data copy correctness. """ assert self.kv_topo is not None + kv_topo = self.kv_topo block_size_ratio = self.block_size // block_size - blocks_data = [] - for i, base_addr in enumerate(self.seen_base_addresses): - # The new block_len is using prefill block_len; - # and num_blocks is multiple with N - kv_block_len = ( - self.get_backend_aware_kv_block_len(layer_idx=i) // block_size_ratio - ) - block_len_per_layer = self.block_len_per_layer[i] // block_size_ratio - num_blocks = self.num_blocks * block_size_ratio - for block_id in range(num_blocks): - block_offset = block_id * block_len_per_layer - addr = base_addr + block_offset - # (addr, len, device id) - blocks_data.append((addr, kv_block_len, self.device_id)) - - if self.kv_topo.is_kv_layout_blocks_first: - # Separate and interleave K/V regions to maintain the same - # descs ordering. This is needed for selecting contiguous heads - # when split across TP ranks. + blocks_data: list[tuple[int, int, int]] = [] + local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank] + + def register_blocks(blocks_data: list[tuple[int, int, int]], mamba: bool): + for i, base_addr in enumerate(local_base_addresses): + # The new block_len is using prefill block_len; + # and num_blocks is multiple with N + kv_block_len = ( + self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=mamba + ) + // block_size_ratio + ) + # Jump one page_size, but ssm page_size may be bigger when kernel + # locks block size to a specific value. + block_len_per_layer = ( + self.block_len_per_layer[i] + // block_size_ratio + * (1 if not mamba else self._physical_blocks_per_logical_kv_block) + ) + num_blocks = self._logical_num_blocks if mamba else self.num_blocks + num_blocks = num_blocks * block_size_ratio for block_id in range(num_blocks): block_offset = block_id * block_len_per_layer addr = base_addr + block_offset - # Register addresses for V cache (K registered first). - v_addr = addr + kv_block_len - blocks_data.append((v_addr, kv_block_len, self.device_id)) - logger.debug( - "Created %s blocks for src engine %s and rank %s on device id %s", - len(blocks_data), - self.engine_id, - self.tp_rank, - self.device_id, - ) + # (addr, len, device id) + blocks_data.append((addr, kv_block_len, self.device_id)) + + if kv_topo.is_kv_layout_blocks_first: + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=mamba + ) + # Separate and interleave K/V regions to maintain the same + # descs ordering. This is needed for selecting contiguous heads + # when split across TP ranks. + for block_id in range(num_blocks): + block_offset = block_id * block_len_per_layer + addr = base_addr + block_offset + # Register addresses for V cache (K registered first). + v_addr = addr + kv_block_len + blocks_data.append((v_addr, second_split, self.device_id)) + logger.debug( + "Created %s blocks for src engine %s and rank %s on device id %s", + len(blocks_data), + self.engine_id, + self.tp_rank, + self.device_id, + ) + + register_blocks(blocks_data, mamba=False) + if self._has_mamba: + assert self.num_descs == len(blocks_data) + logger.debug( + "Registering additional %s local Mamba blocks", len(blocks_data) + ) + register_blocks(blocks_data, mamba=True) descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) # NIXL_INIT_AGENT to be used for preparations of local descs. @@ -1708,7 +1857,8 @@ def add_remote_agent( # local origin:| 0| 1| 8| 12| # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15| assert self.kv_topo is not None - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id) + kv_topo = self.kv_topo + block_size_ratio = kv_topo.block_size_ratio_from_engine_id(engine_id) if engine_id not in self.dst_num_blocks: self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks @@ -1768,48 +1918,86 @@ def add_remote_agent( # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..]. # Register all remote blocks, but only the corresponding kv heads. - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): - # Read our whole local region size from remote. - local_block_len = self.get_backend_aware_kv_block_len(layer_idx=i) - remote_kv_block_len = local_block_len // block_size_ratio - if block_size_ratio > 1: - # using remote kv_block_len as transfer unit - local_block_len = remote_kv_block_len + def register_remote_blocks( + blocks_data: list[tuple[int, int, int]], mamba: bool + ): + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr): + # Read our whole local region size from remote. + local_block_len = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=True, mamba_view=mamba + ) + remote_kv_block_len = local_block_len // block_size_ratio + if block_size_ratio > 1: + # using remote kv_block_len as transfer unit + local_block_len = remote_kv_block_len + + if tp_ratio < 0 and not self.use_mla: + # Remote tp is bigger: read a chunk of local region from remote + local_block_len = local_block_len // (-tp_ratio) + rank_offset = ( + self.tp_rank % tp_ratio * remote_kv_block_len + if indexes_into_remote + else 0 + ) - if tp_ratio < 0 and not self.use_mla: - # Remote tp is bigger: read a chunk of local region from remote - local_block_len = local_block_len // (-tp_ratio) - rank_offset = ( - self.tp_rank % tp_ratio * remote_kv_block_len - if indexes_into_remote - else 0 - ) - for block_id in range(nixl_agent_meta.num_blocks): - block_offset = block_id * nixl_agent_meta.block_lens[i] - # For each block, grab the heads chunk belonging to rank_i - # of size remote_nheads // tp_ratio, which correspond to - # self.block_len == remote_block_len//tp_ratio bytes. - addr = base_addr + block_offset + rank_offset - # (addr, len, device id) - blocks_data.append((addr, local_block_len, nixl_agent_meta.device_id)) - - if self.kv_topo.is_kv_layout_blocks_first: - # With FlashInfer index V separately to allow head splitting. - for block_id in range(nixl_agent_meta.num_blocks): - block_offset = block_id * nixl_agent_meta.block_lens[i] + # Assume same num_blocks for mamba and fa + num_blocks = ( + nixl_agent_meta.num_blocks + if not mamba + else nixl_agent_meta.num_blocks + // self._physical_blocks_per_logical_kv_block + ) + page_size = nixl_agent_meta.block_lens[i] * ( + 1 if not mamba else self._physical_blocks_per_logical_kv_block + ) + for block_id in range(num_blocks): + block_offset = block_id * page_size + # For each block, grab the heads chunk belonging to rank_i + # of size remote_nheads // tp_ratio, which correspond to + # self.block_len == remote_block_len//tp_ratio bytes. addr = base_addr + block_offset + rank_offset - v_addr = addr + nixl_agent_meta.block_lens[i] // 2 + # (addr, len, device id) blocks_data.append( - (v_addr, local_block_len, nixl_agent_meta.device_id) + (addr, local_block_len, nixl_agent_meta.device_id) ) - logger.debug( - "Created %s blocks for dst engine %s with remote rank %s and local rank %s", - len(blocks_data), - engine_id, - remote_tp_rank, - self.tp_rank, - ) + if kv_topo.is_kv_layout_blocks_first: + # With FlashInfer index V separately to allow head splitting. + second_split = self.get_backend_aware_kv_block_len( + layer_idx=i, first_split=False, mamba_view=mamba + ) + # Apply the same scaling as local_block_len above for when we read + # a chunk of local V from `tp_ratio` separate remote workers. + if tp_ratio < 0 and not self.use_mla: + second_split = second_split // (-tp_ratio) + for block_id in range(num_blocks): + block_offset = block_id * page_size + addr = base_addr + block_offset + rank_offset + # Hop over the first split of remote page: either K or Conv. + if mamba: + v_addr = addr + nixl_agent_meta.ssm_sizes[0] + else: + v_addr = addr + nixl_agent_meta.block_lens[i] // 2 + blocks_data.append( + (v_addr, second_split, nixl_agent_meta.device_id) + ) + + logger.debug( + "Created %s blocks for dst engine %s" + " with remote rank %s and local rank %s", + len(blocks_data), + engine_id, + remote_tp_rank, + self.tp_rank, + ) + + register_remote_blocks(blocks_data, mamba=False) + if self._has_mamba: + # Create extra descs for the Mamba "view" of the same KV cache tensors. + logger.debug( + "Registering additional %s remote Mamba blocks", len(blocks_data) + ) + register_remote_blocks(blocks_data, mamba=True) # Register with NIXL. descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type) @@ -1849,6 +2037,9 @@ def _validate_remote_agent_handshake( assert block_size_ratio == 1, ( "HMA does not support different remote block size yet" ) + # Mamba additional constraints + if self._has_mamba: + assert tp_ratio == 1, "Mamba does not support heterogeneous TP yet" kv_cache_layout = ( self.kv_cache_layout @@ -2495,6 +2686,7 @@ def _get_block_descs_ids( A single flattened array is returned for all groups anyway. """ region_ids = np.arange(self.num_regions) + # NOTE (NickLucche) With HMA, every kv group has the same number of layers and # layers from different groups share the same kv tensor. # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be read across all regions, @@ -2505,11 +2697,33 @@ def _get_block_descs_ids( if block_size_ratio is not None: num_blocks = int(num_blocks * block_size_ratio) - # Compute the desc ids for each block. + # Compute desc ids per group using the right stride: FA descs have + # num_blocks entries per region (kernel granularity), SSM descs have + # logical_blocks entries per region (no kernel splitting). region_ids = region_ids[:, None] - block_ids = np.concatenate(block_ids)[None, :] - descs_ids = region_ids * num_blocks + block_ids - return descs_ids.flatten() + if not self._has_mamba: + block_ids = np.concatenate(block_ids)[None, :] + descs_ids = region_ids * num_blocks + block_ids + return descs_ids.flatten() + else: + # NOTE (NickLucche) SSM and Attention blocks regions can be exchanged + # arbitrarily by manager. Therefore, descs are duplicated for SSM and + # Attention like so: + # desc_handle->[descs_fa (all regions) | descs_ssm (all regions)]. + # This is like having two "low-level views" of the same storage. + # `num_fa_descs` offset must be computed per-engine since P and D can + # have different num_blocks (and thus different FA descs counts). + ratio = self._physical_blocks_per_logical_kv_block + # SSM may register fewer num_blocks than FA + logical_blocks = num_blocks // ratio + num_fa_descs = self.num_regions * num_blocks + all_descs = [] + for i, group in enumerate(block_ids): + stride = logical_blocks if self._is_mamba_group[i] else num_blocks + group_arr = np.asarray(group)[None, :] + offset = num_fa_descs if self._is_mamba_group[i] else 0 + all_descs.append((region_ids * stride + group_arr + offset).flatten()) + return np.concatenate(all_descs) def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: """ @@ -2523,16 +2737,22 @@ def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds: block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape( 1, -1 ) + # Mamba blocks have no logical<>physical discrepancy + group_specs = self.kv_cache_config.kv_cache_groups return [ BlockTable.map_to_kernel_blocks( np.array(group), self._physical_blocks_per_logical_kv_block, block_arange, ).tolist() - for group in block_ids + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec) + else group + for i, group in enumerate(block_ids) ] - def get_backend_aware_kv_block_len(self, layer_idx: int) -> int: + def get_backend_aware_kv_block_len( + self, layer_idx: int, first_split: bool = True, mamba_view: bool = False + ) -> int: """ Get the block length for one K/V element (K and V have the same size). @@ -2540,11 +2760,38 @@ def get_backend_aware_kv_block_len(self, layer_idx: int) -> int: block, as K and V are in separate regions. For FlashInfer, this is half the length of the whole block, as K and V share the same region. + Similarly, for SSM-based models, state and conv are interleaved, but crucially + the their size differs. + Reference diagram: + KVCacheTensor (Shared) + / \ + / \ + / \ + Attention (FlashInfer) View Mamba View + | | + | | + +-------------------+ +-------------------+ + | KVCacheTensor | | KVCacheTensor | + | | | | + |<----- page ------>| |<----- page ------->| + | size | | size | + | Key 0 | Val 0 | |Conv 0 | SSM 0 | + | Key 1 | Val 1 | |Conv 1 | SSM 1 | + | ... | ... | | ... | ... | + | Key N-2 | Val N-2 | |Conv N-2| SSM N-2 | + | Key N-1 | Val N-1 | |Conv N-1| SSM N-1 | + +-------------------+ +--------------------+ + |1st_split-2nd_split| |1st_split-2nd_split | """ assert self.kv_topo is not None if self.kv_topo.is_kv_layout_blocks_first: # For indexing only half (either just the K or V part). - block_len = self.block_len_per_layer[layer_idx] // 2 + if mamba_view: + # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so + # we're only transferring the minimum required bytes. + block_len = self._mamba_ssm_size[not first_split] + else: + block_len = self.block_len_per_layer[layer_idx] // 2 else: block_len = self.block_len_per_layer[layer_idx] return block_len diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index 021f0144d81d..4c850fd2f8bd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -24,7 +24,7 @@ ) from vllm.forward_context import ForwardContext from vllm.logger import init_logger -from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import BlockHash @@ -601,7 +601,9 @@ def _register_handlers( def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): layer_names = list(kv_caches.keys()) layers = get_layers_from_vllm_config( - self.spec.vllm_config, Attention, layer_names + self.spec.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] + layer_names, ) attn_backends = { layer_name: layers[layer_name].get_attn_backend() diff --git a/vllm/entrypoints/cli/launch.py b/vllm/entrypoints/cli/launch.py index 6afa2435365a..cc9e467c4604 100644 --- a/vllm/entrypoints/cli/launch.py +++ b/vllm/entrypoints/cli/launch.py @@ -116,6 +116,11 @@ async def run_launch_fastapi(args: argparse.Namespace) -> None: # 2. Build and serve the API server engine_args = AsyncEngineArgs.from_cli_args(args) model_config = engine_args.create_model_config() + + # Render servers preprocess data only — no inference, no quantized kernels. + # Clear quantization so VllmConfig skips quant dtype/capability validation. + model_config.quantization = None + vllm_config = VllmConfig(model_config=model_config) shutdown_task = await build_and_serve_renderer( vllm_config, listen_address, sock, args diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py index 002ae62b8ee8..126e2b4024e8 100644 --- a/vllm/entrypoints/openai/api_server.py +++ b/vllm/entrypoints/openai/api_server.py @@ -29,11 +29,13 @@ from vllm.entrypoints.launcher import serve_http from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args +from vllm.entrypoints.openai.engine.protocol import GenerationError from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.server_utils import ( engine_error_handler, exception_handler, + generation_error_handler, get_uvicorn_log_config, http_exception_handler, lifespan, @@ -263,6 +265,7 @@ def build_app( app.exception_handler(RequestValidationError)(validation_exception_handler) app.exception_handler(EngineGenerateError)(engine_error_handler) app.exception_handler(EngineDeadError)(engine_error_handler) + app.exception_handler(GenerationError)(generation_error_handler) app.exception_handler(Exception)(exception_handler) # Ensure --api-key option from CLI takes precedence over VLLM_API_KEY diff --git a/vllm/entrypoints/openai/server_utils.py b/vllm/entrypoints/openai/server_utils.py index 1453d8083c80..7e9e9a0290e3 100644 --- a/vllm/entrypoints/openai/server_utils.py +++ b/vllm/entrypoints/openai/server_utils.py @@ -21,7 +21,11 @@ from vllm import envs from vllm.engine.protocol import EngineClient from vllm.entrypoints.launcher import terminate_if_errored -from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm.entrypoints.openai.engine.protocol import ( + ErrorInfo, + ErrorResponse, + GenerationError, +) from vllm.entrypoints.utils import create_error_response, sanitize_message from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger @@ -354,6 +358,17 @@ async def engine_error_handler( return JSONResponse(err.model_dump(), status_code=err.error.code) +async def generation_error_handler(req: Request, exc: GenerationError): + """Handle GenerationError without logging stack traces. + + GenerationError is a known, expected error (e.g. KV cache load failure) + that should be returned to the client as a 500 response without polluting + server logs with stack traces. + """ + err = create_error_response(exc) + return JSONResponse(err.model_dump(), status_code=err.error.code) + + async def exception_handler(req: Request, exc: Exception): if req.app.state.args.log_error_stack: logger.exception( diff --git a/vllm/envs.py b/vllm/envs.py index caa2fb38afb6..d6240df36051 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -296,6 +296,16 @@ def use_aot_compile() -> bool: ) +def use_mega_aot_artifact(): + from vllm.utils.torch_utils import is_torch_equal_or_newer + + default_value = ( + "1" if is_torch_equal_or_newer("2.12.0.dev") and use_aot_compile() else "0" + ) + + return os.environ.get("VLLM_USE_MEGA_AOT_ARTIFACT", default_value) == "1" + + def env_with_choices( env_name: str, default: str | None, @@ -616,10 +626,7 @@ def _get_or_set_default() -> str: # Enable loading compiled models directly from cached standalone compile artifacts # without re-splitting graph modules. This reduces overhead during model # loading by using reconstruct_serializable_fn_from_mega_artifact. - "VLLM_USE_MEGA_AOT_ARTIFACT": lambda: os.environ.get( - "VLLM_USE_MEGA_AOT_ARTIFACT", "0" - ) - == "1", + "VLLM_USE_MEGA_AOT_ARTIFACT": use_mega_aot_artifact, # local rank of the process in the distributed setting, used to determine # the GPU device id "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 36ee728dca96..b613f3ba983b 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1142,10 +1142,12 @@ def get_kv_cache_shape( def get_kv_cache_stride_order( include_num_layers_dimension: bool = False, ) -> tuple[int, ...]: - # `stride_order` indicates the permutation that gets - # us from `get_kv_cache_shape` to the actual memory layout we want. - # (num_blocks, num_layers, block_size, head_size) - return (1, 0, 2, 3) if include_num_layers_dimension else (0, 1, 2) + if include_num_layers_dimension: + # MLA kernels require contiguous per-layer KV cache views. + # Identity permutation keeps num_layers first in physical + # layout, signaling cross-layer allocation is unsupported. + return (0, 1, 2, 3) + return (0, 1, 2) @classmethod def get_supported_head_sizes(cls) -> list[int]: diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 51a97e0a2610..534cab1b8276 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/cutlass_moe.py @@ -659,6 +659,13 @@ def run_cutlass_moe_fp4( class CutlassExpertsFp4(mk.FusedMoEExpertsModular): """CUTLASS FP4 fused MoE expert implementation.""" + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Fuse activation scales into w_scale_2 in-place so that + # g1/g2_alphas (which reference the same tensor) stay in sync + # when EPLB rearranges the parameter. + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + @property def expects_unquantized_inputs(self) -> bool: return True diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index 1c86702e9ec1..74096ef6ed6f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -23,6 +23,8 @@ kFp8Dynamic128Sym, kFp8Static128BlockSym, kFp8StaticTensorSym, + kMxfp8Dynamic, + kMxfp8Static, ) from vllm.platforms import current_platform @@ -67,11 +69,54 @@ def _supports_no_act_and_mul() -> bool: """Does not support non-gated MoE (i.e. Nanotron-3-Nano).""" return True + @staticmethod + def _supports_quant_scheme( + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Supports Fp8 per-tensor, Fp8 block, and MXFP8.""" + SUPPORTED_W_A = [ + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kMxfp8Static, kMxfp8Dynamic), + ] + return (weight_key, activation_key) in SUPPORTED_W_A + @staticmethod def _supports_activation(activation: MoEActivation) -> bool: """Supports only SiLU and RELU^2 non-gated activation.""" return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + @staticmethod + def _supports_routing_method( + routing_method: RoutingMethodType, + weight_key: QuantKey | None, + activation_key: QuantKey | None, + ) -> bool: + """Monolithic kernels need to express router support.""" + # NOTE(dbari): TopK routing could also be enabled, but need to validate models + # NOTE(dbari): Default is not implemented and should not be enabled until it is + if (weight_key, activation_key) in [ + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kMxfp8Static, kMxfp8Dynamic), + ]: + # NOTE(rob): potentially allow others here. This is a conservative list. + return routing_method in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym): + # NOTE(dbari): as above, potentially allow others here. + return routing_method in [ + RoutingMethodType.DeepSeekV3, + RoutingMethodType.Llama4, + RoutingMethodType.Renormalize, + RoutingMethodType.RenormalizeNaive, + ] + else: + raise ValueError("Unsupported quantization scheme.") + @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: """Monolithic kernel so only use with naive DP/EP and TP.""" @@ -113,9 +158,10 @@ def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - """Supports Fp8 block.""" + """Supports Fp8 block and MXFP8.""" SUPPORTED_W_A = [ (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kMxfp8Static, kMxfp8Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A @@ -159,6 +205,7 @@ def apply( apply_router_weight_on_input: bool, ): import flashinfer + from flashinfer.fused_moe import Fp8QuantizationType # Pack topk_ids and topk_weights into single tensor # Format: (expert_id << 16) | (weight_bf16.view(int16)) @@ -175,6 +222,16 @@ def apply( assert a1q_scale is not None + is_mxfp8 = self.quant_config.block_shape == [1, 32] + if is_mxfp8: + fp8_quant_type = Fp8QuantizationType.MxFp8 + use_shuffled_weight = True + hidden_states_scale = a1q_scale + else: + fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 + use_shuffled_weight = False + hidden_states_scale = a1q_scale.t().contiguous() + # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the # output tensor in-place so we need to manually copy the result to the # output tensor @@ -183,7 +240,7 @@ def apply( topk_ids=packed_topk_ids, routing_bias=None, hidden_states=hidden_states, - hidden_states_scale=a1q_scale.t().contiguous(), # type: ignore[union-attr] + hidden_states_scale=hidden_states_scale, gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale, gemm2_weights=w2, @@ -197,8 +254,9 @@ def apply( local_num_experts=self.local_num_experts, routed_scaling_factor=None, routing_method_type=1, - use_shuffled_weight=False, + use_shuffled_weight=use_shuffled_weight, weight_layout=0, + fp8_quantization_type=fp8_quant_type, # output=output, ) output.copy_(result) @@ -240,10 +298,11 @@ def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - """Supports Fp8 per-tensor and Fp8 block.""" + """Supports Fp8 per-tensor, Fp8 block, and MXFP8.""" SUPPORTED_W_A = [ (kFp8Static128BlockSym, kFp8Dynamic128Sym), (kFp8StaticTensorSym, kFp8StaticTensorSym), + (kMxfp8Static, kMxfp8Dynamic), ] return (weight_key, activation_key) in SUPPORTED_W_A @@ -256,7 +315,10 @@ def _supports_routing_method( """Monolithic kernels need to express router support.""" # NOTE(dbari): TopK routing could also be enabled, but need to validate models # NOTE(dbari): Default is not implemented and should not be enabled until it is - if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym): + if (weight_key, activation_key) in [ + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + (kMxfp8Static, kMxfp8Dynamic), + ]: # NOTE(rob): potentially allow others here. This is a conservative list. return routing_method in [ RoutingMethodType.DeepSeekV3, @@ -274,7 +336,7 @@ def _supports_routing_method( else: raise ValueError("Unsupported quantization scheme.") - def _apply_per_block( + def _apply_block_scale( self, hidden_states: torch.Tensor, w1: torch.Tensor, @@ -291,32 +353,38 @@ def _apply_per_block( routed_scaling_factor: float | None = None, topk_group: int | None = None, ) -> torch.Tensor: - # Delay import for non-CUDA. import flashinfer + from flashinfer.fused_moe import Fp8QuantizationType assert not apply_router_weight_on_input assert activation == MoEActivation.SILU - - if self.routing_method_type == RoutingMethodType.DeepSeekV3: - router_logits = router_logits.to(torch.float32) - assert self.topk <= global_num_experts assert self.topk <= 10 assert global_num_experts % 4 == 0 - assert self.quant_config.block_shape == [128, 128] - # Routing kernel expects #experts <= #threads 512 + assert self.quant_config.block_shape in [[128, 128], [1, 32]] + # Kernel expects #experts <= #threads 512 assert global_num_experts <= 512 - - # Kernel requires transposed hidden state scales # TODO: fuse into the quant kernel. assert a1q_scale is not None - a1q_scale_t = a1q_scale.t().contiguous() + + if self.routing_method_type == RoutingMethodType.DeepSeekV3: + router_logits = router_logits.to(torch.float32) + + is_mxfp8 = self.quant_config.block_shape == [1, 32] + if is_mxfp8: + fp8_quant_type = Fp8QuantizationType.MxFp8 + use_shuffled_weight = True + hidden_states_scale = a1q_scale + else: + fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 + use_shuffled_weight = False + hidden_states_scale = a1q_scale.t().contiguous() return flashinfer.fused_moe.trtllm_fp8_block_scale_moe( routing_logits=router_logits, routing_bias=e_score_correction_bias, hidden_states=hidden_states, - hidden_states_scale=a1q_scale_t, + hidden_states_scale=hidden_states_scale, gemm1_weights=w1, gemm1_weights_scale=self.quant_config.w1_scale, gemm2_weights=w2, @@ -330,7 +398,8 @@ def _apply_per_block( local_num_experts=self.local_num_experts, routed_scaling_factor=routed_scaling_factor, routing_method_type=self.routing_method_type, - use_shuffled_weight=False, + use_shuffled_weight=use_shuffled_weight, + fp8_quantization_type=fp8_quant_type, ) def _apply_per_tensor( @@ -409,7 +478,7 @@ def apply( topk_group: int | None = None, ) -> torch.Tensor: if self.quant_config.block_shape is not None: - return self._apply_per_block( + return self._apply_block_scale( hidden_states, w1, w2, @@ -441,6 +510,6 @@ def apply( ) else: raise NotImplementedError( - "Only per-block and per-tensor quantization are supported in " - f"{self.__class__.__name__}." + "Only per-block, per-tensor, and MXFP8 quantization are " + f"supported in {self.__class__.__name__}." ) diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py index 174c581b396f..87b1eb9fd58d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py @@ -56,10 +56,25 @@ def __init__( # g1_scale_c = a13_scale * w13_scale_2 / a2_scale self.g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale else: - self.g1_scale_c = ( - torch.ones_like(self.quant_config.a1_gscale) - * self.quant_config.a2_gscale - ) + self.g1_scale_c = self.quant_config.a2_gscale.clone() + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + # Recompute g1_scale_c since g1_alphas was just fused in-place. + # Register as a layer parameter so EPLB rearranges it alongside + # other expert weights. + assert self.quant_config.g1_alphas is not None + assert self.quant_config.a2_gscale is not None + if self.moe_config.is_act_and_mul: + g1_scale_c = self.quant_config.g1_alphas * self.quant_config.a2_gscale + else: + g1_scale_c = self.quant_config.a2_gscale.clone() + layer.register_parameter( + "g1_scale_c", + torch.nn.Parameter(g1_scale_c, requires_grad=False), + ) + self.g1_scale_c = layer.g1_scale_c @staticmethod def _supports_current_device() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py index fb8a18ef33e5..5805a4dd5bf6 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py @@ -49,6 +49,10 @@ def __init__( ) self.out_dtype = moe_config.in_dtype + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.BatchedExperts diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index e58d52eee980..91f7a83f6fce 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -61,6 +61,11 @@ def is_valid_flashinfer_cutlass_fused_moe( class FlashInferExperts(mk.FusedMoEExpertsModular): + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self.quant_config.use_nvfp4_w4a4: + layer.w13_weight_scale_2.data.mul_(layer.w13_input_scale) + layer.w2_weight_scale_2.data.mul_(layer.w2_input_scale) + def __init__( self, moe_config: mk.FusedMoEConfig, diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 7135cbbd2d7c..75283b9bbe39 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -1421,19 +1421,23 @@ def _maybe_make_contiguous( weights = list(self.named_parameters()) weights = [(name, _maybe_make_contiguous(name, p)) for name, p in weights] + # `w13_input_scale` and `w2_input_scale` are global per-tensor + # activation scales shared across all experts (e.g. NVFP4). + # They are broadcast views (stride 0) from .expand() and are + # not actual expert weights, so exclude them from EPLB. + NON_EXPERT_WEIGHTS = { + "e_score_correction_bias", + "w13_input_scale", + "w2_input_scale", + } + assert all( weight.is_contiguous() for name, weight in weights if not (name.startswith("_shared_experts.") or name.startswith("_gate.")) + and name not in NON_EXPERT_WEIGHTS ) - # Filter out the non-expert weights. - # `e_score_correction_bias` is a bias for each logical expert, - # with shape (num_logical_experts,), not an expert weight. - NON_EXPERT_WEIGHTS = { - "e_score_correction_bias", - } - return [ weight.view(self.local_num_experts, -1) for name, weight in weights diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 7100c87c91c7..a6b498834017 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -489,6 +489,9 @@ def __init__( self.max_num_tokens = max_num_tokens self.num_dispatchers = num_dispatchers + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # noqa: B027 + pass + @staticmethod def is_monolithic() -> bool: raise NotImplementedError("Implemented by subclasses.") diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 48ca03f666c5..a63c02663886 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -444,7 +444,7 @@ def convert_to_fp8_moe_kernel_format( Fp8MoeBackend.FLASHINFER_CUTLASS, Fp8MoeBackend.FLASHINFER_TRTLLM, ]: - w13, w2, w13_scale = prepare_fp8_moe_layer_for_fi( + w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_fi( layer=layer, w13=w13, w2=w2, @@ -512,6 +512,21 @@ def make_fp8_moe_quant_config( g1_alphas=(w1_scale * a1_scale).squeeze(), g2_alphas=(w2_scale * a2_scale).squeeze(), ) + # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to + # _mxfp8_e4m3_quantize rather than standard FP8 block quantization. + # Non-swizzled layout is required since the TRTLLM kernel expects + # scales in (num_tokens, hidden_dim // 32) format. + if block_shape == [1, 32]: + return FusedMoEQuantConfig.make( + "mxfp8", + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + is_nvfp4_scale_swizzled=False, + ) + # All other backends use normal config. return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py index 49406ba935e2..ed3af4b5a474 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp8.py @@ -1,44 +1,87 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from enum import Enum +import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + backend_to_kernel_cls, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + kMxfp8Dynamic, + kMxfp8Static, +) logger = init_logger(__name__) +_SUPPORTED_BACKENDS: frozenset[Fp8MoeBackend] = frozenset( + { + Fp8MoeBackend.FLASHINFER_TRTLLM, + } +) -class MxFp8MoeBackend(Enum): - FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" +_BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = { + "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM, +} + + +def _select_kernel_cls( + backend: Fp8MoeBackend, + config: FusedMoEConfig, +) -> type[mk.FusedMoEExperts]: + """Select the first supported expert class for the MXFP8 config.""" + activation_format = ( + mk.FusedMoEActivationFormat.BatchedExperts + if config.moe_parallel_config.use_batched_activation_format + else mk.FusedMoEActivationFormat.Standard + ) + last_reason: str | None = None + for cls in backend_to_kernel_cls(backend): + supported, reason = cls.is_supported_config( + cls, + config, + kMxfp8Static, + kMxfp8Dynamic, + activation_format, + ) + if supported: + return cls + last_reason = reason + raise ValueError( + f"No supported MXFP8 expert class for {backend.value}: {last_reason}" + ) def select_mxfp8_moe_backend( config: FusedMoEConfig, -) -> MxFp8MoeBackend: +) -> tuple[Fp8MoeBackend, type[mk.FusedMoEExperts]]: + """Select the MXFP8 MoE backend and the best expert class. + + Returns: + A tuple of (fp8_backend, experts_cls). + """ if config.is_lora_enabled: raise NotImplementedError("LoRA is not supported for MXFP8 MoE.") - AVAILABLE_BACKENDS = [ - MxFp8MoeBackend.FLASHINFER_TRTLLM, - ] - runner_backend = config.moe_backend if runner_backend != "auto": - mapping = { - "flashinfer_trtllm": MxFp8MoeBackend.FLASHINFER_TRTLLM, - } - if backend := mapping.get(runner_backend): - logger.info_once( - "Using '%s' MxFp8 MoE backend (user-requested).", - backend.value, + backend = _BACKEND_NAME_MAP.get(runner_backend) + if backend is None: + raise ValueError( + f"moe_backend='{runner_backend}' is not supported for " + f"MXFP8 MoE. Expected one of " + f"{list(_BACKEND_NAME_MAP.keys())}." ) - return backend - raise ValueError( - f"moe_backend='{runner_backend}' is not supported for MXFP8 MoE. " - f"Expected one of {list(mapping.keys())}." + logger.info_once( + "Using '%s' MxFp8 MoE backend (user-requested).", + backend.value, ) + return backend, _select_kernel_cls(backend, config) + + # Auto-select: pick the first supported backend. + for backend in _SUPPORTED_BACKENDS: + logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) + return backend, _select_kernel_cls(backend, config) - # Auto-select: only one backend available for now. - backend = AVAILABLE_BACKENDS[0] - logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value) - return backend + raise ValueError("No MXFP8 MoE backends available.") diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index b06cf49cfd81..8a224cb39e7c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -374,11 +374,13 @@ def make_nvfp4_moe_quant_config( w2_scale=w2_scale, ) - g1_alphas = a13_scale * w13_scale_2 - g2_alphas = a2_scale * w2_scale_2 + # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas. + # The expert's process_weights_after_loading will fuse activation + # scales in-place. Since the quant config references the same tensor + # as the registered parameter, EPLB rearrangement stays in sync. return nvfp4_moe_quant_config( - g1_alphas=g1_alphas, - g2_alphas=g2_alphas, + g1_alphas=w13_scale_2, + g2_alphas=w2_scale_2, a1_gscale=(1.0 / a13_scale), a2_gscale=(1.0 / a2_scale), w1_scale=w13_scale, diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 019e408c1959..4adb7f1cfa0e 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -199,7 +199,7 @@ def _mxfp8_e4m3_quantize( ) -> tuple[torch.Tensor, torch.Tensor]: assert A_scale is None assert not per_act_token_quant - assert block_shape is None + assert block_shape is None or block_shape == [1, 32] return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout) diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 2fb54e7751a0..e08a6456aba7 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -31,6 +31,7 @@ "torchao", "inc", "mxfp4", + "mxfp8", "petit_nvfp4", "cpu_awq", ] @@ -129,6 +130,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: ) from .moe_wna16 import MoeWNA16Config from .mxfp4 import Mxfp4Config + from .mxfp8 import Mxfp8Config from .petit import PetitNvFp4Config from .ptpc_fp8 import PTPCFp8Config from .torchao import TorchAOConfig @@ -156,6 +158,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "auto-round": INCConfig, "inc": INCConfig, "mxfp4": Mxfp4Config, + "mxfp8": Mxfp8Config, "petit_nvfp4": PetitNvFp4Config, "cpu_awq": CPUAWQConfig, } diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index f35a4c0b977c..29115fbbc255 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -570,6 +570,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: shared_experts=layer.shared_experts, routing_tables=layer._maybe_init_expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def maybe_make_prepare_finalize( self, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 977612313f63..78644f74d288 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -25,13 +25,13 @@ FusedMoeWeightScaleSupported, ) from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, convert_to_fp8_moe_kernel_format, make_fp8_moe_kernel, make_fp8_moe_quant_config, select_fp8_moe_backend, ) from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( - MxFp8MoeBackend, select_mxfp8_moe_backend, ) from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( @@ -1394,6 +1394,7 @@ def process_weights_after_loading(self, layer: FusedMoE) -> None: shared_experts=layer.shared_experts, routing_tables=layer._maybe_init_expert_routing_tables(), ) + self.moe_kernel.fused_experts.process_weights_after_loading(layer) def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig: return make_nvfp4_moe_quant_config( @@ -1711,8 +1712,7 @@ def __init__( self.quant_config = quant_config assert self.quant_config.is_checkpoint_mxfp8_serialized - # Select MXFP8 MoE backend - self.mxfp8_backend = select_mxfp8_moe_backend(self.moe) + self.mxfp8_backend, _ = select_mxfp8_moe_backend(self.moe) def create_weights( self, @@ -1942,7 +1942,7 @@ def get_fused_moe_quant_config( @property def is_monolithic(self) -> bool: - return self.mxfp8_backend == MxFp8MoeBackend.FLASHINFER_TRTLLM + return self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM def apply_monolithic( self, @@ -1955,7 +1955,7 @@ def apply_monolithic( Fp8QuantizationType, ) - assert self.mxfp8_backend == MxFp8MoeBackend.FLASHINFER_TRTLLM + assert self.mxfp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM if layer.enable_eplb: raise NotImplementedError( diff --git a/vllm/model_executor/layers/quantization/mxfp8.py b/vllm/model_executor/layers/quantization/mxfp8.py new file mode 100644 index 000000000000..5b4564bea31c --- /dev/null +++ b/vllm/model_executor/layers/quantization/mxfp8.py @@ -0,0 +1,354 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Online MXFP8 (microscaling FP8, block-32) quantization config and methods.""" + +from typing import Any + +import torch +from torch.nn import Module + +from vllm.logger import init_logger +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod +from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import ( + select_mxfp8_moe_backend, +) +from vllm.model_executor.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import ( + Fp8Config, + Fp8KVCacheMethod, + Fp8OnlineLinearMethod, + Fp8OnlineMoEMethod, + _copy_missing_attrs, +) +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( + MXFP8_BLOCK_SIZE, + Mxfp8LinearBackend, + Mxfp8LinearOp, + mxfp8_e4m3_quantize, + swizzle_mxfp8_scale, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.model_loader.weight_utils import ( + initialize_single_dummy_weight, +) +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class Mxfp8Config(Fp8Config): + """Config class for online MXFP8 MoE quantization.""" + + def __init__( + self, + activation_scheme: str = "dynamic", + ignored_layers: list[str] | None = None, + ) -> None: + if activation_scheme != "dynamic": + raise ValueError("mxfp8 only supports dynamic activation scheme.") + super().__init__( + is_checkpoint_fp8_serialized=False, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=None, + ) + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "mxfp8" + + @classmethod + def get_min_capability(cls) -> int: + return 100 + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "Mxfp8Config": + activation_scheme = cls.get_from_keys_or( + config, ["activation_scheme"], "dynamic" + ) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls( + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + skip_with_substr=True, + ): + return UnquantizedLinearMethod() + return Mxfp8OnlineLinearMethod(self) + elif isinstance(layer, FusedMoE): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + skip_with_substr=True, + ): + return UnquantizedFusedMoEMethod(layer.moe_config) + return Mxfp8OnlineMoEMethod(self, layer) + elif isinstance(layer, Attention): + return Fp8KVCacheMethod(self) + return None + + +class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod): + """Online MXFP8 linear method. + Loads bf16/fp16 checkpoints and quantizes weights to MXFP8 (microscaling + FP8 with block-32 scales) during weight loading. + + Args: + quant_config: The MXFP8 quantization config. + """ + + uses_meta_device: bool = True + + def __init__(self, quant_config: "Mxfp8Config"): + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + self.mxfp8_linear = Mxfp8LinearOp(self._select_backend()) + logger.info_once( + "Using %s backend for MXFP8 GEMM", self.mxfp8_linear.backend.value + ) + + @staticmethod + def _select_backend() -> Mxfp8LinearBackend: + try: + from vllm.utils import flashinfer as fi + + _ = fi.mm_mxfp8 + return Mxfp8LinearBackend.FLASHINFER_CUTLASS + except Exception: + logger.warning( + "FlashInfer mm_mxfp8 not available, " + "falling back to MXFP8 emulation backend." + ) + return Mxfp8LinearBackend.EMULATION + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + if input_size_per_partition % MXFP8_BLOCK_SIZE != 0: + raise ValueError( + f"MXFP8 requires input_size_per_partition " + f"({input_size_per_partition}) to be divisible by " + f"{MXFP8_BLOCK_SIZE}." + ) + + super().create_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + weight_fp8, weight_scale = mxfp8_e4m3_quantize(layer.weight.contiguous()) + + if self.mxfp8_linear.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS: + N, K = layer.weight.shape[0], layer.weight.shape[1] + weight_scale = swizzle_mxfp8_scale(weight_scale, N, K) + + layer.input_scale = None + replace_parameter(layer, "weight", weight_fp8.data) + replace_parameter(layer, "weight_scale", weight_scale.data) + + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.mxfp8_linear.apply( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + out_dtype=self.out_dtype, + bias=bias, + ) + + +class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod): + """MoE method for online MXFP8 (block) quantization.""" + + uses_meta_device: bool = True + + def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module): + FusedMoEMethodBase.__init__(self, layer.moe_config) + self.quant_config = quant_config + assert not quant_config.is_checkpoint_fp8_serialized + assert quant_config.activation_scheme == "dynamic" + + self.weight_block_size = [1, MXFP8_BLOCK_SIZE] + self.block_quant = True + self.weight_scale_name = "weight_scale" + + self.fp8_backend, self.experts_cls = select_mxfp8_moe_backend(config=self.moe) + + def create_weights( + self, + layer: Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + if ( + hidden_size % MXFP8_BLOCK_SIZE != 0 + or intermediate_size_per_partition % MXFP8_BLOCK_SIZE != 0 + ): + raise ValueError( + "Online MXFP8 MoE requires hidden/intermediate sizes divisible " + f"by {MXFP8_BLOCK_SIZE}." + ) + + super().create_weights( + layer=layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + params_dtype=params_dtype, + **extra_weight_attrs, + ) + + w13_weight_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // MXFP8_BLOCK_SIZE, + dtype=torch.uint8, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.zeros( + num_experts, + hidden_size, + intermediate_size_per_partition // MXFP8_BLOCK_SIZE, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + layer.weight_block_size = [1, MXFP8_BLOCK_SIZE] + + def _quantize_mxfp8_moe_weight( + self, weight: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Batch quantization: bf16/fp16 weights -> MXFP8 (fp8 + uint8 scales).""" + num_batches = weight.size(0) + w_quant = [] + w_scales = [] + for i in range(num_batches): + mx_fp8_quant, mx_fp8_scale = mxfp8_e4m3_quantize( + weight[i], is_sf_swizzled_layout=False + ) + w_quant.append(mx_fp8_quant) + w_scales.append(mx_fp8_scale) + + return torch.stack(w_quant), torch.stack(w_scales) + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + if layer.w13_weight.device == torch.device("meta"): + w13_weight = torch.nn.Parameter( + torch.empty_like(layer.w13_weight, device=layer._load_device), + requires_grad=False, + ) + set_weight_attrs( + w13_weight, {"weight_loader": layer.w13_weight.weight_loader} + ) + _copy_missing_attrs(layer.w13_weight, w13_weight) + layer.register_parameter("w13_weight", w13_weight) + initialize_single_dummy_weight(layer.w13_weight) + if layer.w2_weight.device == torch.device("meta"): + w2_weight = torch.nn.Parameter( + torch.empty_like(layer.w2_weight, device=layer._load_device), + requires_grad=False, + ) + set_weight_attrs( + w2_weight, {"weight_loader": layer.w2_weight.weight_loader} + ) + _copy_missing_attrs(layer.w2_weight, w2_weight) + layer.register_parameter("w2_weight", w2_weight) + initialize_single_dummy_weight(layer.w2_weight) + + fp8_dtype = current_platform.fp8_dtype() + w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype) + w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype) + w13_scale = layer.w13_weight_scale + w2_scale = layer.w2_weight_scale + + w13, w13_scale = self._quantize_mxfp8_moe_weight(layer.w13_weight) + w2, w2_scale = self._quantize_mxfp8_moe_weight(layer.w2_weight) + + self._setup_kernel( + layer, + w13, + w2, + w13_scale, + w2_scale, + layer.w13_input_scale, + layer.w2_input_scale, + ) + + layer._already_called_process_weights_after_loading = True diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 42677a5927b3..66300ceaefab 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -267,16 +267,6 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( num_experts=w13.size(0), is_gated_activation=is_gated, ) - - # We do not need to make this a parameter, because - # it is not used during the weight (re)-loading process. - if is_gated: - layer.g1_scale_c = a13_scale * w13_scale_2 / a2_scale - else: - layer.g1_scale_c = torch.ones_like(a13_scale) / a2_scale - layer.a1_gscale = 1.0 / a13_scale - layer.g1_alphas = a13_scale * w13_scale_2 - layer.g2_alphas = a2_scale * w2_scale_2 else: # Swizzle the block scales for other FI NVFP4 MoE kernels. w13_scale = swizzle_blockscale(w13_scale) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py index 322b3a6e86b1..271bcf168386 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_utils.py @@ -305,6 +305,81 @@ def align_fp8_moe_weights_for_fi( return padded_w13, padded_w2, padded_intermediate +def _shuffle_mxfp8_moe_weights( + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + is_gated: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Preprocess MXFP8 weights and scales for the FlashInfer TRT-LLM kernel. + + Following flashinfer/tests/moe/test_trtllm_gen_fused_moe.py: + 1. reorder_rows_for_gated_act_gemm (interleave gate/up rows) + 2. shuffle_matrix_a (weight data layout shuffle) + 3. shuffle_matrix_sf_a (scale factor layout shuffle) + """ + from flashinfer import ( + reorder_rows_for_gated_act_gemm, + shuffle_matrix_a, + shuffle_matrix_sf_a, + ) + + epilogue_tile_m = 128 + num_experts = w13.shape[0] + intermediate_size = w13.shape[1] // 2 + hidden_size = w13.shape[2] + + w13_interleaved: list[torch.Tensor] = [] + w13_scale_interleaved: list[torch.Tensor] = [] + for i in range(num_experts): + if is_gated: + w13_interleaved.append( + reorder_rows_for_gated_act_gemm( + w13[i].reshape(2 * intermediate_size, -1) + ) + ) + w13_scale_interleaved.append( + reorder_rows_for_gated_act_gemm( + w13_scale[i].reshape(2 * intermediate_size, -1) + ) + ) + else: + w13_interleaved.append(w13[i]) + w13_scale_interleaved.append(w13_scale[i]) + + w13_shuffled: list[torch.Tensor] = [] + w2_shuffled: list[torch.Tensor] = [] + w13_scale_shuffled: list[torch.Tensor] = [] + w2_scale_shuffled: list[torch.Tensor] = [] + for i in range(num_experts): + w13_shuffled.append( + shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m) + ) + w2_shuffled.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)) + w13_scale_shuffled.append( + shuffle_matrix_sf_a( + w13_scale_interleaved[i] + .view(torch.uint8) + .reshape(2 * intermediate_size, -1), + epilogue_tile_m, + ) + ) + w2_scale_shuffled.append( + shuffle_matrix_sf_a( + w2_scale[i].view(torch.uint8).reshape(hidden_size, -1), + epilogue_tile_m, + ) + ) + + w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn) + w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn) + w13_scale_out = torch.stack(w13_scale_shuffled).reshape(w13_scale.shape) + w2_scale_out = torch.stack(w2_scale_shuffled).reshape(w2_scale.shape) + + return w13_out, w2_out, w13_scale_out, w2_scale_out + + def prepare_fp8_moe_layer_for_fi( layer: torch.nn.Module, w13: torch.Tensor, @@ -314,7 +389,7 @@ def prepare_fp8_moe_layer_for_fi( w2_scale: torch.Tensor, w2_input_scale: torch.Tensor | None, is_trtllm: bool = False, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """ Convert Fp8 MoE weights to flashinfer kernel format @@ -329,10 +404,33 @@ def prepare_fp8_moe_layer_for_fi( block_quant = ( hasattr(layer, "weight_block_size") and layer.weight_block_size is not None ) + is_mxfp8 = block_quant and w13_scale.dtype == torch.uint8 + is_gated = layer.activation.is_gated + + # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle. + if is_mxfp8 and is_trtllm: + # FlashInfer TRT-LLM SwiGLU expects [up; gate] but vLLM stores + # [gate; up]. Swap both weights and scales before interleaving. + if layer.moe_config.is_act_and_mul: + w13 = swap_w13_to_w31(w13) + # Scales may be 2D [E, flat] from _quantize_mxfp8_moe_weight; + # reshape to 3D so swap_w13_to_w31 can flip the two halves, + # then flatten back. + if w13_scale.ndim == 2: + num_rows = w13.shape[1] # 2 * intermediate_size + w13_scale = w13_scale.reshape(w13_scale.shape[0], num_rows, -1) + w13_scale = swap_w13_to_w31(w13_scale) + w13_scale = w13_scale.reshape(w13_scale.shape[0], -1) + else: + w13_scale = swap_w13_to_w31(w13_scale) + + w13, w2, w13_scale, w2_scale = _shuffle_mxfp8_moe_weights( + w13, w2, w13_scale, w2_scale, is_gated + ) + return w13, w2, w13_scale, w2_scale # Some FI MoE kernels require internal alignment of 16 # for the gate-up proj. Pad the weights to respect this. - is_gated = layer.activation.is_gated if not block_quant: min_alignment = 16 if is_gated else 128 w13, w2, new_intermediate = align_fp8_moe_weights_for_fi( @@ -369,4 +467,4 @@ def prepare_fp8_moe_layer_for_fi( w13_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE) w2_scale.clamp_(min=_FI_CUTLASS_MIN_BLOCK_SCALE) - return w13, w2, w13_scale + return w13, w2, w13_scale, w2_scale diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 12a1799d157c..1170a2d3a77c 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -149,6 +149,12 @@ def __str__(self): kStatic128BlockScale = ScaleDesc(torch.float32, True, GroupShape(128, 128)) kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True) +kMxfp8StaticScale = ScaleDesc(torch.uint8, True, GroupShape(1, 32)) +kMxfp8Static = QuantKey(FP8_DTYPE, kMxfp8StaticScale, symmetric=True) + +kMxfp8DynamicScale = ScaleDesc(torch.uint8, False, GroupShape(1, 32)) +kMxfp8Dynamic = QuantKey(FP8_DTYPE, kMxfp8DynamicScale, symmetric=True) + kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64)) kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True) diff --git a/vllm/model_executor/models/mistral_large_3_eagle.py b/vllm/model_executor/models/mistral_large_3_eagle.py index 4567f24fdade..3fcc048f9fa9 100644 --- a/vllm/model_executor/models/mistral_large_3_eagle.py +++ b/vllm/model_executor/models/mistral_large_3_eagle.py @@ -74,6 +74,7 @@ def __init__( prefix=maybe_prefix(prefix, "fc"), ) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.aux_hidden_state_layers: tuple[int, ...] = () self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) diff --git a/vllm/model_executor/models/mllama4.py b/vllm/model_executor/models/mllama4.py index da9836a95e0b..a36b1fa57082 100644 --- a/vllm/model_executor/models/mllama4.py +++ b/vllm/model_executor/models/mllama4.py @@ -38,7 +38,6 @@ from vllm.config import VllmConfig, set_current_vllm_config from vllm.config.multimodal import BaseDummyOptions from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention import MMEncoderAttention from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.model_executor.layers.linear import ( @@ -872,10 +871,7 @@ def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings: if image_input is None: return [] - with ( - set_forward_context(None, self.vllm_config), - ): - return self._process_image_input(image_input) + return self._process_image_input(image_input) def forward( self, diff --git a/vllm/model_executor/models/nano_nemotron_vl.py b/vllm/model_executor/models/nano_nemotron_vl.py index b32067557622..93cbb417ed51 100644 --- a/vllm/model_executor/models/nano_nemotron_vl.py +++ b/vllm/model_executor/models/nano_nemotron_vl.py @@ -105,15 +105,16 @@ class NanoNemotronVLAudioFeatureInputs(TensorSchema): """ Dimensions: - - b: Number of audio clips + - c: Number of audio clips (possibly flattened across audio items) + - b: Number of original audio items - t: Audio feature length - f: Feature size (mel bins) """ type: Literal["audio_features"] = "audio_features" - input_audio_features: Annotated[torch.Tensor, TensorShape("b", "t", "f")] - feature_attention_mask: Annotated[torch.Tensor, TensorShape("b", "t")] - audio_feature_lengths: Annotated[torch.Tensor, TensorShape("b")] + input_audio_features: Annotated[torch.Tensor, TensorShape("c", "t", "f")] + feature_attention_mask: Annotated[torch.Tensor, TensorShape("c", "t")] + audio_num_clips: list[int] MAX_AUDIO_LEN_S = 10 * 60 # 10 minutes @@ -292,19 +293,110 @@ def image_to_pixel_values( return pixel_values +def _compute_aspect_preserving_size( + orig_w: int, + orig_h: int, + target_num_patches: int, + patch_size: int, + downsample_ratio: float, +) -> tuple[int, int]: + """Compute target pixel dimensions that preserve aspect ratio. + + Mirrors Megatron-LM image_processing.py video frame resizing: + target area in patch-grid space is *target_num_patches*, distributed + according to the source aspect ratio, then snapped to a multiple of + the required divisor (2 for pixel-shuffle). + """ + aspect_wh = orig_w / max(orig_h, 1) + ph = round(math.sqrt(target_num_patches / aspect_wh)) + pw = round(math.sqrt(target_num_patches * aspect_wh)) + ph = max(ph, 1) + pw = max(pw, 1) + + reduction_factor = int(round(1 / downsample_ratio)) + required_divisor = reduction_factor # 2 for pixel-shuffle + if required_divisor > 1: + rem_h = ph % required_divisor + rem_w = pw % required_divisor + ph_up = ph + (required_divisor - rem_h if rem_h else 0) + ph_down = ph - rem_h + pw_up = pw + (required_divisor - rem_w if rem_w else 0) + pw_down = pw - rem_w + if ph_up * pw_up <= target_num_patches: + ph, pw = ph_up, pw_up + else: + ph = max(required_divisor, ph_down) + pw = max(required_divisor, pw_down) + + return pw * patch_size, ph * patch_size # (width, height) in pixels + + +def _get_video_target_size_and_feature_size( + orig_w: int, + orig_h: int, + target_patches: int, + maintain_aspect_ratio: bool, + patch_size: int, + downsample_ratio: float, +) -> tuple[int, int, int]: + """Compute target (width, height) and feature_size for video resize and token count. + + Used by video_to_pixel_values (resize) and get_video_replacement_internvl + (seq length calc) so both use the same dimensions. + """ + if maintain_aspect_ratio: + target_w, target_h = _compute_aspect_preserving_size( + orig_w=orig_w, + orig_h=orig_h, + target_num_patches=target_patches, + patch_size=patch_size, + downsample_ratio=downsample_ratio, + ) + else: + reduction_factor = int(round(1 / downsample_ratio)) + side = int(math.sqrt(target_patches)) + side = max(reduction_factor, (side // reduction_factor) * reduction_factor) + target_w = side * patch_size + target_h = side * patch_size + + feature_size = int((target_h // patch_size) * downsample_ratio) * int( + (target_w // patch_size) * downsample_ratio + ) + return target_w, target_h, feature_size + + def video_to_pixel_values( video: npt.NDArray, *, input_size: int, - max_num_tiles: int = 1, - use_thumbnail: bool, + video_target_num_patches: int | None = None, + video_maintain_aspect_ratio: bool = False, + patch_size: int = 16, + downsample_ratio: float = 0.5, ) -> torch.Tensor: - assert max_num_tiles == 1, "Video modality always uses one tile" - # (num_frames, H, W, C) -> (num_frames, C, H, W) video_tensor = torch.from_numpy(video).permute(0, 3, 1, 2) - if video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size: + if video_target_num_patches is not None: + # Resize to target patch count (aspect-preserving or square). + orig_h, orig_w = video_tensor.shape[2], video_tensor.shape[3] + target_w, target_h, _ = _get_video_target_size_and_feature_size( + orig_w=orig_w, + orig_h=orig_h, + target_patches=video_target_num_patches, + maintain_aspect_ratio=video_maintain_aspect_ratio, + patch_size=patch_size, + downsample_ratio=downsample_ratio, + ) + if video_tensor.shape[2] != target_h or video_tensor.shape[3] != target_w: + video_tensor = torch.nn.functional.interpolate( + video_tensor, + size=(target_h, target_w), + mode="bicubic", + align_corners=False, + antialias=True, + ) + elif video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size: video_tensor = torch.nn.functional.interpolate( video_tensor, size=(input_size, input_size), @@ -809,9 +901,9 @@ def _preprocess_image( "which should be a single string" ) parts = [x for x in re.split(r"()", text[0]) if x] - assert parts.count("") == len(pixel_values_lst), ( - "the number of tokens in the text should be the " - "same as the number of images" + assert parts.count("") == len(num_tokens_per_image), ( + f"Expected {len(num_tokens_per_image)} tokens in text " + f"but found {parts.count('')}" ) for i, (feature_size, num_patches) in enumerate( @@ -868,6 +960,33 @@ def __init__( self.video_token = video_token self.video_pruning_rate = video_pruning_rate + # Video params live exclusively in vision_config + vision_config = getattr(config, "vision_config", config) + self.video_temporal_patch_size: int = getattr( + vision_config, "video_temporal_patch_size", 1 + ) + self.video_maintain_aspect_ratio: bool = getattr( + vision_config, "video_maintain_aspect_ratio", False + ) + + # Resolve video frame target size: exactly one of video_target_num_patches + # or video_target_img_size may be set (mirrors Megatron's + # DynamicResolutionImageTilingStrategy validation). + target_num_patches = getattr(vision_config, "video_target_num_patches", None) + target_img_size = getattr(vision_config, "video_target_img_size", None) + if target_num_patches is not None and target_img_size is not None: + raise ValueError( + "Exactly one of video_target_num_patches or " + "video_target_img_size must be set, got both" + ) + if target_num_patches is not None: + self.video_target_num_patches: int | None = target_num_patches + elif target_img_size is not None: + base_patches = math.ceil(target_img_size / config.patch_size) + self.video_target_num_patches = base_patches * base_patches + else: + self.video_target_num_patches = None + self.audio_extractor: ParakeetExtractor | None = None raw_sound_config = getattr(config, "sound_config", None) if raw_sound_config is not None: @@ -883,6 +1002,27 @@ def __init__( IMG_CONTEXT, add_special_tokens=False ) + @cached_property + def num_video_token(self) -> int: + """Token count per video frame, accounting for video_target_num_patches. + + When video_target_num_patches is set the per-frame feature count + differs from the image-based num_image_token. We use a square + dummy (1:1) to compute the feature_size because the dummy video is + square and the user confirmed that is acceptable. + """ + if self.video_target_num_patches is not None: + _, _, feature_size = _get_video_target_size_and_feature_size( + orig_w=self.image_size, + orig_h=self.image_size, + target_patches=self.video_target_num_patches, + maintain_aspect_ratio=self.video_maintain_aspect_ratio, + patch_size=self.config.patch_size, + downsample_ratio=self.config.downsample_ratio, + ) + return feature_size + return self.num_image_token + @property def supports_video(self) -> bool: return self.video_token_id is not None @@ -900,14 +1040,15 @@ def image_token_id(self) -> int: def _videos_to_pixel_values_lst( self, videos: list[npt.NDArray], - max_num_tiles: int, ) -> list[torch.Tensor]: return [ video_to_pixel_values( video, input_size=self.image_size, - max_num_tiles=max_num_tiles, - use_thumbnail=self.use_thumbnail, + video_target_num_patches=self.video_target_num_patches, + video_maintain_aspect_ratio=self.video_maintain_aspect_ratio, + patch_size=self.config.patch_size, + downsample_ratio=self.config.downsample_ratio, ) for video in videos ] @@ -916,7 +1057,6 @@ def _preprocess_video( self, text: list[str], videos: list[tuple[npt.NDArray, dict[str, Any]]], - max_num_tiles: int, ): if len(videos) == 0 or not self.supports_video: video_inputs = {} @@ -925,7 +1065,6 @@ def _preprocess_video( video_metadata_lst = [v[1] for v in videos] pixel_values_lst_video = self._videos_to_pixel_values_lst( videos_lst, - max_num_tiles=max_num_tiles, ) # We use frame duration in milliseconds (as integer) to ensure @@ -952,12 +1091,10 @@ def _preprocess_video( "frame_duration_ms": torch.tensor(frame_duration_ms_lst), } - image_size: int = self.config.force_image_size patch_size: int = self.config.patch_size downsample_ratio = self.config.downsample_ratio - tokens_in_single_frame = int( - (image_size * image_size // patch_size**2) * (downsample_ratio**2) - ) + + T = self.video_temporal_patch_size for pixel_values, video_metadata, frames_indices, frame_duration_ms in zip( pixel_values_lst_video, @@ -966,6 +1103,11 @@ def _preprocess_video( frame_duration_ms_lst, ): num_frames = pixel_values.shape[0] + frame_h, frame_w = pixel_values.shape[-2], pixel_values.shape[-1] + tokens_in_single_frame = int( + (frame_h * frame_w // patch_size**2) * (downsample_ratio**2) + ) + num_tubelets = math.ceil(num_frames / T) if T > 1 else num_frames if ( self.video_pruning_rate is not None @@ -974,18 +1116,18 @@ def _preprocess_video( # Start of EVS-specific code num_tokens = compute_retained_tokens_count( tokens_per_frame=tokens_in_single_frame, - num_frames=num_frames, + num_frames=num_tubelets, q=self.video_pruning_rate, ) # Here we just need placeholders that won't actually be replaced - # we just need to make sure the total number of tokens is correct # assign all tokens to the first frame - tokens_per_frame = [num_tokens] + [0] * (num_frames - 1) + tokens_per_frame = [num_tokens] + [0] * (num_tubelets - 1) # End of EVS-specific code else: - tokens_per_frame = [tokens_in_single_frame] * num_frames + tokens_per_frame = [tokens_in_single_frame] * num_tubelets video_repl = self.get_video_repl( tokens_per_frame=tokens_per_frame, @@ -995,6 +1137,7 @@ def _preprocess_video( img_start_token_ids=self._img_start_token_ids, img_end_token_ids=self._img_end_token_ids, img_context_token_ids=self._img_context_token_ids, + video_temporal_patch_size=T, ) # video_repl.full is a list of token IDs @@ -1012,7 +1155,7 @@ def _preprocess_audio( audios: list[npt.NDArray], ): if len(audios) == 0: - return text, {} + return text, {"audio_num_clips": []} assert self.audio_extractor is not None extractor = self.audio_extractor @@ -1036,13 +1179,10 @@ def _preprocess_audio( sampling_rate=extractor.sampling_rate, return_tensors="pt", ) - input_audio_features = audio_inputs.input_features - feature_attention_mask = audio_inputs.attention_mask - audio_feature_lengths = feature_attention_mask.sum(dim=1) audio_inputs = { - "input_audio_features": input_audio_features, - "feature_attention_mask": feature_attention_mask, - "audio_feature_lengths": audio_feature_lengths, + "input_audio_features": audio_inputs.input_features, + "feature_attention_mask": audio_inputs.attention_mask, + "audio_num_clips": audio_inputs.audio_num_clips, } return text, audio_inputs @@ -1073,7 +1213,6 @@ def __call__( text, video_inputs = self._preprocess_video( text=text, videos=videos, - max_num_tiles=1, ) text, audio_inputs = self._preprocess_audio( @@ -1127,6 +1266,7 @@ def get_video_repl( img_start_token_ids: list[int], img_end_token_ids: list[int], img_context_token_ids: list[int], + video_temporal_patch_size: int = 1, ) -> PromptUpdateDetails[list[int]]: """ Build prompt replacement for a video. @@ -1146,31 +1286,60 @@ def get_video_repl( - EVS real (called from get_real_video_repl_for_evs) - different value per frame Args: tokens_per_frame (list[int]): number of tokens per frame - frames_indices (list[int]): frame indices + (one per tubelet when T > 1) + frames_indices (list[int]): orig. frame indices + (one per frame, before tubelet subsampling) frame_duration_ms (int): duration of each frame in milliseconds tokenizer (TokenizerLike): tokenizer to use for tokenizing frame separators img_start_token_ids (list[int]): pre-tokenized IMG_START tokens img_end_token_ids (list[int]): pre-tokenized IMG_END tokens img_context_token_ids (list[int]): pre-tokenized IMG_CONTEXT tokens + video_temporal_patch_size (int): temporal patch size for videos """ # TODO: Add support of frame_duration_ms to be None # At preprocessing step we should allow absent / metadata without # frames_indices field. timestamps_enabled = frame_duration_ms is not None - - if timestamps_enabled: + T = video_temporal_patch_size + num_frames = len(frames_indices) + + if T > 1 and timestamps_enabled: + all_timestamps = calculate_timestamps(frames_indices, frame_duration_ms) + + frame_separators = [] + for group_idx, i in enumerate(range(0, num_frames, T)): + group_frames = [] + for j in range(T): # Every frame in the group + frame_idx = i + j + if frame_idx < num_frames: + # Valid idx (haven't padded to mult. of T yet) + ts = all_timestamps[frame_idx] + frame_str = "Frame" if j == 0 else "frame" + group_frames.append( + f"{frame_str} {frame_idx + 1} sampled at {ts:.2f} seconds" + ) + if group_frames: + # Join by `and` if there are >1 frame, otherwise no `and` + # Prepend \n to match training format (except first group) + sep = " and ".join(group_frames) + ": " + if group_idx > 0: + sep = "\n" + sep + frame_separators.append(sep) + elif timestamps_enabled: timestamps = calculate_timestamps(frames_indices, frame_duration_ms) assert len(timestamps) == len(tokens_per_frame), ( "timestamps and tokens_per_frame must have the same length" ) frame_separators = [ - f"Frame {i + 1} sampled at {timestamp:.2f} seconds: " + ("\n" if i > 0 else "") + + f"Frame {i + 1} sampled at {timestamp:.2f} seconds: " for i, timestamp in enumerate(timestamps) ] else: frame_separators = [ - f"Frame {i + 1}: " for i, _ in enumerate(tokens_per_frame) + ("\n" if i > 0 else "") + f"Frame {i + 1}: " + for i, _ in enumerate(tokens_per_frame) ] # Tokenize frame separator independently @@ -1295,10 +1464,13 @@ def get_num_frames_with_most_features( max_videos = mm_counts.get("video", 0) processor = self.get_hf_processor() # we get the CustomProcessor here + T = processor.video_temporal_patch_size max_image_tokens = self.get_max_image_tokens() * max_images - max_total_frames = (seq_len - max_image_tokens) // processor.num_image_token - max_frames_per_video = max_total_frames // max(max_videos, 1) + tokens_per_tubelet = processor.num_video_token + max_total_tubelets = (seq_len - max_image_tokens) // tokens_per_tubelet + max_tubelets_per_video = max_total_tubelets // max(max_videos, 1) + max_frames_per_video = max_tubelets_per_video * T return max(max_frames_per_video, 1) def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor: @@ -1467,6 +1639,18 @@ def apply( processor_inputs.hf_processor_mm_kwargs = hf_processor_mm_kwargs + # Prepend "This is a video:\n" before