diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 9da81d0d15f..65df467acf3 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -155,6 +155,17 @@ steps: commands: - pytest -v -s models/inkling/rocm +- label: ":nvidia: (B200) Flint" + key: flint-unit-tests-b200 + timeout_in_minutes: 30 + device: b200-k8s + source_file_dependencies: + - vllm/models/flint/ + - tests/models/flint/ + commands: + - uv pip install --system 'gguf>=0.19.0' + - pytest -v -s models/flint + - label: ":nvidia: (B200) Kimi K3" key: kimi-k3-unit-tests-b200 timeout_in_minutes: 40 diff --git a/docker/Dockerfile b/docker/Dockerfile index cd6d86ff74d..c359a93c378 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -972,12 +972,15 @@ RUN set -eux; \ # Install vllm-openai dependencies (saves ~2.6s per build) # These are stable packages that don't depend on vLLM itself -# From versions.json: .openai_server_extras.timm, .openai_server_extras.runai_model_streamer +# From versions.json: .openai_server_extras.timm, .openai_server_extras.runai_model_streamer, +# .openai_server_extras.gguf (NVFP4 GGUF checkpoints, e.g. RecursiveSI/flint-nvfp4-gguf) ARG TIMM_VERSION=">=1.0.17" ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7" +ARG GGUF_VERSION=">=0.19.0" RUN --mount=type=cache,target=/opt/uv/cache \ uv pip install --system accelerate 'modelscope<1.38' \ - "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}" + "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}" \ + "gguf${GGUF_VERSION}" # ============================================================ # VLLM INSTALLATION (depends on build stage) diff --git a/docker/versions.json b/docker/versions.json index 12ec728ec8f..75c21afafc6 100644 --- a/docker/versions.json +++ b/docker/versions.json @@ -84,6 +84,9 @@ }, "RUNAI_MODEL_STREAMER_VERSION": { "default": ">=0.15.7" + }, + "GGUF_VERSION": { + "default": ">=0.19.0" } } } diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9e1e3d1e7bd..2265193e3d8 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -400,6 +400,7 @@ th { | `FalconForCausalLM` | Falcon | `tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc. | | ✅︎ | | `FalconMambaForCausalLM` | FalconMamba | `tiiuae/falcon-mamba-7b`, `tiiuae/falcon-mamba-7b-instruct`, etc. | | ✅︎ | | `FalconH1ForCausalLM` | Falcon-H1 | `tiiuae/Falcon-H1-34B-Base`, `tiiuae/Falcon-H1-34B-Instruct`, etc. | ✅︎ | ✅︎ | +| `FlintForCausalLM` | Flint | `RecursiveSI/flint`, `RecursiveSI/flint-nvfp4-gguf` (GGUF) | | ✅︎ | | `GemmaForCausalLM` | Gemma | `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc. | ✅︎ | ✅︎ | | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ | | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | diff --git a/recipes/models/RecursiveSI/flint.yaml b/recipes/models/RecursiveSI/flint.yaml new file mode 100644 index 00000000000..9a117293206 --- /dev/null +++ b/recipes/models/RecursiveSI/flint.yaml @@ -0,0 +1,152 @@ +# yaml-language-server: $schema=https://recipes.vllm.ai/recipe.schema.json + +meta: + title: "Flint" + slug: "flint" + provider: "RecursiveSI" + description: "39.8B-parameter (2.5B active) MoE with hybrid sliding-window / full attention, reasoning and tool calling; BF16 and NVFP4 checkpoints on a single Blackwell GPU." + date_added: 2026-10-18 + date_updated: 2026-10-21 + tasks: + - text + performance_headline: "Single GB200, NVFP4: ~2.0 ms per output token at batch 1 with 100K context" + related_recipes: [] + default_hardware: gb200 + hardware: + gb200: verified + +model: + model_id: "RecursiveSI/flint" + min_vllm_version: "nightly" + install: + pip: + note: "Use a nightly wheel. The NVFP4 GGUF checkpoint also needs `pip install 'gguf>=0.19.0'`." + docker: + note: "Use the vLLM nightly image." + architecture: moe + parameter_count: "39.8B" + active_parameters: "2.5B" + context_length: 262144 + base_args: [] + base_env: {} + rust_frontend: unsupported + +features: + tool_calling: + description: "Automatic tool choice with the GLM-4.7 XML tool-call parser (Flint's chat template emits //); tools are dropped from the prompt when tool_choice is none, as in the reference runtime." + args: ["--enable-auto-tool-choice", "--tool-call-parser", "glm47", "--exclude-tools-when-tool-choice-none"] + reasoning: + description: "Return the … block in the reasoning field; the model decides per turn whether to reason." + args: ["--reasoning-parser", "glm47"] + +opt_in_features: [] + +variants: + default: + precision: bf16 + vram_minimum_gb: 96 + description: "BF16 safetensors (79.6 GB of weights)." + nvfp4: + model_id: "RecursiveSI/flint-nvfp4-gguf" + precision: nvfp4 + vram_minimum_gb: 29 + description: "NVFP4 GGUF served weight-only (W4A16); 23.5 GB resident after load (the LM head is dequantized to BF16). Config and tokenizer come from the BF16 repo." + extra_args: ["--hf-config-path", "RecursiveSI/flint", "--tokenizer", "RecursiveSI/flint"] + +compatible_strategies: + - single_node_tp + +hardware_overrides: {} + +strategy_overrides: + single_node_tp: + tp: 1 + +guide: | + ## Overview + + Flint is a 39.8B-parameter mixture-of-experts model (256 experts, top-6 sigmoid routing plus a + shared expert; ~2.5B parameters active per token) with 49 layers: 35 sliding-window layers + (window 513, RoPE) and 14 full-attention layers (no RoPE), GQA 32/4, attention sinks, and a + 262,144-token context. Two checkpoints are served: + + - [RecursiveSI/flint](https://huggingface.co/RecursiveSI/flint): BF16 safetensors. + - [RecursiveSI/flint-nvfp4-gguf](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf): a single + NVFP4 GGUF (repacked losslessly at load and run as W4A16, like the reference runtime). It ships + no config or tokenizer; both come from the BF16 repo. + + Tested end to end on one GB200 (TP=1, BF16 KV cache, CUDA graphs, no torch.compile) for both + checkpoints. On a single GPU the Flint kernels are always used (fused MoE, small-batch CuTe GEMV, + CuTe decode attention, low-latency decode path for <= 8 requests); TP > 1 and expert parallelism + fall back to stock vLLM layers and are not validated. + + ## Launch + + BF16: + + ```bash + vllm serve RecursiveSI/flint --served-model-name flint \ + --enable-auto-tool-choice --tool-call-parser glm47 --reasoning-parser glm47 \ + --exclude-tools-when-tool-choice-none + ``` + + NVFP4 (requires `gguf>=0.19.0`): + + ```bash + vllm serve RecursiveSI/flint-nvfp4-gguf \ + --hf-config-path RecursiveSI/flint --tokenizer RecursiveSI/flint \ + --served-model-name flint \ + --enable-auto-tool-choice --tool-call-parser glm47 --reasoning-parser glm47 \ + --exclude-tools-when-tool-choice-none + ``` + + ## Client usage + + Use `/v1/chat/completions`. Sampling defaults from `generation_config.json`: temperature 1.0, + top_p 0.95, top_k 20 (greedy decoding is not validated). Reasoning is returned in `reasoning` + and counts toward `max_tokens`; use `max_tokens` of at least 32768 for hard problems. + + ```bash + curl localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d '{ + "model": "flint", + "messages": [{"role": "user", "content": "What is 17*23? Answer briefly."}], + "max_tokens": 1024 + }' + ``` + + Accuracy on GB200 with this recipe (BF16 / NVFP4): GSM8K 0.950 / 0.950, GPQA-Diamond (4 repeats) + 0.721 / 0.696, MRCR (48 samples, up to 114K-token prompts) ~0.47 / ~0.46, tool-calling suite + (19 cases x 10: typed args, tool selection, parallel calls, multi-turn, tool_choice none, + escaping, streaming) 10/10 on every case. + + ## Benchmarking + + Forward time on one GB200 (`vllm bench forward`, batch x context, us): + + | shape | BF16 | NVFP4 | + |---|---|---| + | decode 1 x 100K | 2341 | 1783 | + | decode 8 x 100K | 6836 | 5703 | + | decode 32 x 100K | 19580 | 16716 | + | prefill 8K @ 100K | 145617 | 144765 | + + Serving (NVFP4, `vllm bench serve`, random 1000 in / 1000 out): TPOT 2.0 ms at concurrency 1, + 2.8 ms at 8, 4.9 ms at 32, 8.5 ms at 128; median inter-token latency 2.06 ms at batch 1 with + 100K context. Prefill runs eagerly and is CPU-bound on kernel-launch overhead: TTFT is ~60 ms + for a short prompt and ~125 ms for an 8K-token prompt. + + ## Troubleshooting + + - Rust frontend: not supported for this recipe. It rejects `--exclude-tools-when-tool-choice-none`, + and for NVFP4 also `--hf-config-path` / `--tokenizer`; keep the default Python frontend. + - Reasoning parser: the reference runtime's `deepseek-v3` name maps to vLLM's `deepseek_v3` + parser, which only extracts reasoning when the request enables `thinking`; with Flint the + reasoning then stays in `content`. Use `glm47`. + - If `content` contains literal `` or `` text, the server was started without + the parsers above. + - Batch-1 decode time varies by ~3-4% between sessions (a GPU-side state, not configuration). + + ## References + + - [Flint (BF16) model card](https://huggingface.co/RecursiveSI/flint) + - [Flint NVFP4 model card](https://huggingface.co/RecursiveSI/flint-nvfp4-gguf) diff --git a/tests/models/flint/test_decode_attention.py b/tests/models/flint/test_decode_attention.py new file mode 100644 index 00000000000..143c2c907c7 --- /dev/null +++ b/tests/models/flint/test_decode_attention.py @@ -0,0 +1,278 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The Flint split-KV decode attention kernel must match an fp64 reference at +bf16 output-rounding level, mirror FlashInfer trtllm-gen's numerics, and stay +correct under CUDA graph replay with new sequence lengths.""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cuda() or not current_platform.is_device_capability_family( + 100 +): + pytest.skip("Flint decode attention targets SM100", allow_module_level=True) + +from flashinfer.decode import trtllm_batch_decode_with_kv_cache # noqa: E402 + +from vllm.models.flint.ops.decode_attention import flint_decode_attention # noqa: E402 + +HQ, HKV, D, PAGE = 32, 4, 128, 16 +SCALE = D**-0.5 +WINDOW = 513 + + +def _make( + seq_lens: list[int], seed: int = 0 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Random q/K/V/sinks; blocks are shuffled so block tables are non-trivial.""" + g = torch.Generator(device="cuda").manual_seed(seed) + pages = [(s + PAGE - 1) // PAGE for s in seq_lens] + num_blocks = sum(pages) + 5 + kv = torch.randn(num_blocks, HKV, PAGE, 2 * D, device="cuda", generator=g) + perm = torch.randperm(num_blocks, device="cuda", generator=g).int() + bt = torch.zeros(len(seq_lens), max(pages + [1]) + 2, dtype=torch.int32) + bt = bt.cuda() + o = 0 + for i, n in enumerate(pages): + bt[i, :n] = perm[o : o + n] + o += n + q = torch.randn(len(seq_lens), HQ, D, device="cuda", generator=g) * 2 + sinks = torch.randn(HQ, device="cuda", generator=g) * 2 + sl = torch.tensor(seq_lens, dtype=torch.int32, device="cuda") + return q.bfloat16(), kv.bfloat16(), bt, sl, sinks + + +def _reference(q, kv, bt, seq_lens, sinks, window): + out = torch.zeros(q.shape, dtype=torch.float64, device=q.device) + for i, s in enumerate(seq_lens.tolist()): + if s == 0: + continue + n = (s + PAGE - 1) // PAGE + x = kv[bt[i, :n].long()].permute(1, 0, 2, 3).reshape(HKV, n * PAGE, 2 * D) + x = x[:, max(s - window, 0) if window else 0 : s].double() + k, v = x[..., :D], x[..., D:] + scores = torch.einsum("hgd,htd->hgt", q[i].double().view(HKV, -1, D), k) + scores = scores * SCALE + sink = sinks.double().view(HKV, -1, 1) + m = torch.maximum(scores.amax(-1, keepdim=True), sink) + p = torch.exp(scores - m) + denom = p.sum(-1, keepdim=True) + torch.exp(sink - m) + out[i] = (torch.einsum("hgt,htd->hgd", p, v) / denom).reshape(HQ, D) + return out + + +def _flashinfer(q, kv, bt, seq_lens, sinks, window): + workspace = torch.zeros(128 << 20, dtype=torch.uint8, device="cuda") + k_cache, v_cache = kv.split(D, dim=-1) + return trtllm_batch_decode_with_kv_cache( + q, + (k_cache, v_cache), + workspace, + bt, + seq_lens, + int(seq_lens.max()), + bmm1_scale=SCALE, + window_left=window - 1 if window else -1, + sinks=sinks, + kv_layout="HND", + backend="trtllm-gen", + ) + + +def _flint(q, kv, bt, seq_lens, sinks, window, out=None): + out = torch.empty_like(q) if out is None else out + return flint_decode_attention(q, kv, bt, seq_lens, sinks, window, SCALE, out) + + +def _rel_err(x, ref): + """|x - ref| relative to each output row's max magnitude.""" + scale = ref.abs().amax(-1, keepdim=True).clamp_min(1e-6) + return (x.double() - ref).abs() / scale + + +@pytest.mark.parametrize( + "seq_lens", + [ + [1], + [17], + [1000], + [20011], + [5, 77, 1000, 4097, 16, 33, 513, 3001], + [100, 2000, 37] * 10 + [1, 9000], + ], + ids=["1", "17", "1000", "20011", "mixed8", "mixed32"], +) +@pytest.mark.parametrize("window", [None, WINDOW], ids=["full", "swa"]) +@torch.inference_mode() +def test_matches_reference_and_flashinfer(seq_lens: list[int], window: int | None): + args = _make(seq_lens) + ref = _reference(*args, window) + out = _flint(*args, window) + fi_err = _rel_err(_flashinfer(*args, window), ref) + err = _rel_err(out, ref) + # One bf16 ulp of the row max is 2^-8; trtllm-gen also rounds its split + # partials to bf16, so its error bounds ours. + assert err.max() < 8e-3 + assert err.max() <= fi_err.max() + 1e-3 + assert err.mean() <= fi_err.mean() + 1e-4 + + +@pytest.mark.parametrize("seq_len", [513, 514, 529, 530, 545, 20000]) +@torch.inference_mode() +def test_window_crosses_block_boundary(seq_len: int): + """The first attended page is partial; earlier tokens must be masked.""" + args = _make([seq_len, seq_len + 7], seed=seq_len) + err = _rel_err(_flint(*args, WINDOW), _reference(*args, WINDOW)) + assert err.max() < 8e-3 + + +@torch.inference_mode() +def test_launch_from_another_stream_waits_for_the_previous_one(): + """Launches share the split workspace and arrival counters, so a launch + from a different stream must not start before the previous one ends.""" + args = _make([20011] * 4) + _flint(*args, None) + stream_a, stream_b = torch.cuda.Stream(), torch.cuda.Stream() + torch.accelerator.synchronize() + with torch.cuda.stream(stream_a): + torch.cuda._sleep(100_000_000) # ~50 ms at ~2 GHz + a_done = torch.cuda.Event() + _flint(*args, None) + a_done.record() + with torch.cuda.stream(stream_b): + _flint(*args, None) + b_done = torch.cuda.Event() + b_done.record() + b_done.synchronize() + assert a_done.query() + + +@torch.inference_mode() +def test_empty_request_writes_zeros(): + q, kv, bt, sl, sinks = _make([0, 40]) + out = torch.full_like(q, float("nan")) + _flint(q, kv, bt, sl, sinks, None, out) + assert (out[0] == 0).all() + assert _rel_err(out[1:], _reference(q, kv, bt, sl, sinks, None)[1:]).max() < 8e-3 + + +@torch.inference_mode() +def test_probabilities_summed_in_fp32(): + """Like trtllm-gen, P is rounded to bf16 for P @ V but the softmax + denominator sums the fp32 probabilities. With V = 1, one max-logit token + per page (so every split sees the same running max) and all other + probabilities just below a bf16 rounding midpoint, the output is visibly + below 1 (it would be exactly 1 with fp32 P or a bf16-P denominator).""" + n = 1000 + q, kv, bt, sl, sinks = _make([n]) + kv.zero_() + q.zero_() + q[..., 0] = 1.0 + q[..., 1] = 2**-6 + flat = kv.view(-1, HKV, PAGE, 2 * D)[bt[0, : (n + PAGE - 1) // PAGE].long()] + flat[..., 0] = -7.8125 # p = exp2((s - 0) * scale * log2e) = 0.50195.. + flat[..., 1] = 0.9296875 + flat[:, :, 0, :2] = 0.0 # first token of every page at the max (p = 1) + flat[..., D:] = 1.0 + kv.view(-1, HKV, PAGE, 2 * D)[bt[0, : (n + PAGE - 1) // PAGE].long()] = flat + sinks.fill_(-100.0) + out = _flint(q, kv, bt, sl, sinks, None) + expected = _flashinfer(q, kv, bt, sl, sinks, None) + assert (expected == 0.99609375).all() + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + +@torch.inference_mode() +def test_cuda_graph_replay_with_new_lengths(): + """Grid shape depends only on the batch size; the split-combine counters + reset themselves, so replays with different lengths stay correct.""" + lens_a = [30000, 17, 4100, 515] + lens_b = [1, 31000, 513, 9000] + q, kv, bt, sl, sinks = _make(lens_b) + out = torch.empty_like(q) + _, _, bt_a, sl_a, _ = _make(lens_a) + bt_a = torch.nn.functional.pad(bt_a, (0, bt.shape[1] - bt_a.shape[1])) + static_bt, static_sl = bt_a.clone(), sl_a.clone() + for window in (None, WINDOW): + _flint(q, kv, static_bt, static_sl, sinks, window, out) # warm up / compile + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _flint(q, kv, static_bt, static_sl, sinks, window, out) + for bt_x, sl_x in ((bt, sl), (bt_a, sl_a), (bt, sl)): + static_bt.copy_(bt_x) + static_sl.copy_(sl_x) + graph.replay() + ref = _reference(q, kv, static_bt, static_sl, sinks, window) + assert _rel_err(out, ref).max() < 8e-3 + + +def _early_trigger_writer(): + """A PDL kernel that releases its dependents immediately, idles ~2 us, + then writes one token's K/V into the cache, like + ``qk_norm_rope_cache(launch_dependents=True)``.""" + import cutlass + import cutlass.cute as cute + from cutlass import Int32, Int64 + + class Writer: + @cute.jit + def __call__(self, src, dst, cycles: Int32, stream): + self.kernel(src, dst, cycles).launch( + grid=(1, 1, 1), block=(256, 1, 1), stream=stream, use_pdl=True + ) + + @cute.kernel + def kernel(self, src, dst, cycles: Int32): + cute.arch.griddepcontrol_wait() + cute.arch.griddepcontrol_launch_dependents() + t0 = cute.arch.clock64() + while cute.arch.clock64() - t0 < Int64(cycles): + pass + tid, _, _ = cute.arch.thread_idx() + dst[tid] = src[tid] + + from cutlass.cute.runtime import make_fake_stream + from quack.compile_utils import make_fake_tensor + + bf16 = cutlass.BFloat16 + return cute.compile( + Writer(), + make_fake_tensor(bf16, (256,), divisibility=8), + make_fake_tensor(bf16, (256,), divisibility=8), + Int32(1), + make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +@pytest.mark.parametrize("window", [None, WINDOW], ids=["full", "swa"]) +@torch.inference_mode() +def test_current_token_written_by_pdl_predecessor(window: int | None): + """K/V pages are prefetched before the grid-dependency wait, except the + page holding the current token: K/V written there by an early-triggering + predecessor must be seen.""" + writer = _early_trigger_writer() + # The last write (request 3, head 3) races the attention kernel; with one + # token, a stale read would change that output entirely. + lens = [20011, 4096, 17, 1] + q, kv, bt, sl, sinks = _make(lens) + new_kv = torch.randn(len(lens), HKV, 2 * D, device="cuda").bfloat16() * 4 + out = torch.empty_like(q) + blks = [int(bt[i, (s - 1) // PAGE]) for i, s in enumerate(lens)] + _flint(q, kv, bt, sl, sinks, window, out) # compile / workspace + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for i, s in enumerate(lens): + for h in range(HKV): + writer(new_kv[i, h], kv[blks[i], h, (s - 1) % PAGE], 4000) + _flint(q, kv, bt, sl, sinks, window, out) + kv_before = kv.clone() + graph.replay() + torch.accelerator.synchronize() + for i, s in enumerate(lens): + kv_before[blks[i], :, (s - 1) % PAGE] = new_kv[i] + assert torch.equal(kv, kv_before) + ref = _reference(q, kv, bt, sl, sinks, window) + assert _rel_err(out, ref).max() < 8e-3 diff --git a/tests/models/flint/test_flint_moe.py b/tests/models/flint/test_flint_moe.py new file mode 100644 index 00000000000..58d34f1abc4 --- /dev/null +++ b/tests/models/flint/test_flint_moe.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint's fused NVFP4 W4A16 MoE (vllm/models/flint/ops/moe.py) must follow the +reference rounding rule on both the decode (GEMV) and the grouped-GEMM paths +and route exactly like sigmoid top-k.""" + +import pytest +import torch + +from vllm.models.flint.gguf_nvfp4 import dequantize_nvfp4 +from vllm.models.flint.ops import moe +from vllm.platforms import current_platform + +E, H, INTER, TOPK = 256, 2048, 512, 6 + +pytestmark = pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="SM100 only", +) + + +def _nvfp4(rows: int, cols: int, gen: torch.Generator): + w = torch.randint( + 0, 256, (E, rows, cols // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + # E4M3 scales over ~2^-6..2^2 with all mantissas. + exp = torch.randint(1, 10, (E, rows, cols // 16), device="cuda", generator=gen) + man = torch.randint(0, 8, (E, rows, cols // 16), device="cuda", generator=gen) + s = ((exp << 3) | man).to(torch.uint8).view(torch.float8_e4m3fn) + gs = torch.rand(E, device="cuda", generator=gen) * 2e-3 + 1e-3 + return w, s, gs + + +@pytest.fixture(scope="module") +def weights(): + gen = torch.Generator(device="cuda").manual_seed(0) + w13, s13, gs13 = _nvfp4(2 * INTER, H, gen) + w2, s2, gs2 = _nvfp4(H, INTER, gen) + router = (torch.randn(E, H, device="cuda", generator=gen) * 0.02).bfloat16() + return dict(w13=w13, s13=s13, gs13=gs13, w2=w2, s2=s2, gs2=gs2, router=router) + + +@pytest.fixture(scope="module") +def flint_weights(weights): + return moe.FlintMoEWeights.from_modelopt( + weights["w13"], + weights["s13"], + weights["gs13"], + weights["w2"], + weights["s2"], + weights["gs2"], + ) + + +def _inputs(num_tokens: int, seed: int | None = None): + gen = torch.Generator(device="cuda").manual_seed( + num_tokens if seed is None else seed + ) + x = torch.randn(num_tokens, H, device="cuda", generator=gen) + x = (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True))).bfloat16() + shared = (torch.randn(num_tokens, H, device="cuda", generator=gen) * 0.1).bfloat16() + return x, shared + + +def _reference(x, w, ids, tw, shared, bf16_points: bool): + """MoE on exactly dequantized weights with fp32 GEMMs. + + bf16_points=False: everything in fp32 (returns fp32). + bf16_points=True: the reference (TRTLLM MoE) rounding rule (returns bf16): + fp32 gate/up, act = bf16(silu(g) * u), y_e = bf16(acc_e * global_scale_e), + routed = bf16(sum_e tw_e * y_e) in fp32 (tw already bf16-rounded by the + router), out = bf16(routed + shared). + """ + bf = torch.bfloat16 + routed = torch.zeros(x.shape[0], H, device=x.device) + for e in ids.unique().tolist(): + tok, slot = (ids == e).nonzero(as_tuple=True) + w13 = dequantize_nvfp4(w["w13"][e], w["s13"][e], torch.tensor(1.0)) + w2 = dequantize_nvfp4(w["w2"][e], w["s2"][e], torch.tensor(1.0)) + gu = (x[tok].float() @ w13.T) * w["gs13"][e] + g, u = gu[:, :INTER], gu[:, INTER:] + if bf16_points: + h = (torch.nn.functional.silu(g) * u).to(bf).float() + y = ((h @ w2.T) * w["gs2"][e]).to(bf).float() + else: + h = torch.nn.functional.silu(g) * u + y = (h @ w2.T) * w["gs2"][e] + routed.index_add_(0, tok, tw[tok, slot, None] * y) + if not bf16_points: + return routed + shared.float() + return (routed.to(bf).float() + shared.float()).to(bf) + + +def _rel(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).norm() / b.float().norm()) + + +def _run(x, router, fw, shared, decode_max: int | None = None): + """flint_moe with an optionally forced decode/grouped dispatch.""" + old = moe.DECODE_MAX_TOKENS + if decode_max is not None: + moe.DECODE_MAX_TOKENS = decode_max + try: + counter = moe.make_route_counter("cuda") + return moe.flint_moe(x, router, fw, TOPK, counter, shared) + finally: + moe.DECODE_MAX_TOKENS = old + + +@torch.inference_mode() +def test_dequant_exact(): + """The E2M1 * E4M3 decode is exact for every (code, scale) pair.""" + import triton + import triton.language as tl + + @triton.jit + def kernel(w_ptr, s_ptr, out_ptr, N: tl.constexpr): + offs = tl.arange(0, N) + o = moe._dequant8(tl.load(w_ptr + offs), tl.load(s_ptr + offs)) + for j in tl.static_range(8): + tl.store(out_ptr + offs * 8 + j, o[j]) + + # Every non-NaN E4M3 scale applied to one 16-code block holding codes + # 0..15 (two int32 words). + scales = torch.arange(256, dtype=torch.uint8, device="cuda") + scales = scales[(scales & 0x7F) != 0x7F] + codes = torch.arange(16, dtype=torch.uint8, device="cuda") + block = codes[0::2] | (codes[1::2] << 4) + words = block.repeat(scales.numel()).view(torch.int32) + n = words.numel() + padded_w = torch.zeros(512, dtype=torch.int32, device="cuda") + padded_s = torch.zeros(512, dtype=torch.uint8, device="cuda") + padded_w[:n] = words + padded_s[:n] = scales.repeat_interleave(2) + out_dtype = torch.bfloat16 if moe._DEQUANT_BF16 else torch.float16 + padded_out = torch.empty(512 * 8, dtype=out_dtype, device="cuda") + kernel[(1,)](padded_w, padded_s, padded_out, N=512) + ref = dequantize_nvfp4( + block.repeat(scales.numel()).view(-1, 8), + scales.view(torch.float8_e4m3fn).view(-1, 1), + torch.tensor(1.0), + ) + torch.testing.assert_close( + padded_out[: n * 8].float(), ref.flatten(), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize( + "num_tokens", + [1, 5, moe.ROUTE_GEMV_MAX_TOKENS, moe.ROUTE_GEMV_MAX_TOKENS + 1, 300], +) +@torch.inference_mode() +def test_routing(weights, num_tokens: int): + x, _ = _inputs(num_tokens) + counter = moe.make_route_counter("cuda") + ids, tw = moe.route(x, weights["router"], TOPK, counter) + if num_tokens > moe.ROUTE_GEMV_MAX_TOKENS: + # Same cuBLAS bf16 GEMM as the kernel path; only top-k is under test. + logits = (x @ weights["router"].T).float() + else: + logits = (x.float() @ weights["router"].float().T).bfloat16().float() + ref_s, ref_ids = torch.sigmoid(logits).topk(TOPK, dim=-1) + torch.testing.assert_close( + ids.sort(-1).values, ref_ids.int().sort(-1).values, rtol=0, atol=0 + ) + # Renormalized as the reference does, then rounded to bf16. + ref_tw = ref_s / (ref_s.sum(-1, keepdim=True) + 1e-20) + ref_tw = ref_tw.bfloat16().float() + order, ref_order = ids.argsort(-1), ref_ids.argsort(-1) + got, ref_tw = tw.gather(1, order), ref_tw.gather(1, ref_order) + # The fp32 sigmoid / sum may differ in the last bit: rare 1-ulp bf16 flips. + assert (got == ref_tw).float().mean() > 0.99 + torch.testing.assert_close(got, ref_tw, rtol=2**-7, atol=0) + + +@torch.inference_mode() +def test_topk_ties_pick_lowest_expert(): + logits = torch.zeros(2, E, device="cuda", dtype=torch.bfloat16) + logits[0, [7, 3, 200, 9, 100, 50, 2]] = 1.0 + logits[1, :] = -1.0 + ids = torch.empty(2, TOPK, dtype=torch.int32, device="cuda") + tw = torch.empty(2, TOPK, dtype=torch.float32, device="cuda") + moe._topk_kernel[(2,)](logits, ids, tw, E=E, TOPK=TOPK) + assert ids[0].tolist() == [2, 3, 7, 9, 50, 100] + assert ids[1].tolist() == [0, 1, 2, 3, 4, 5] + torch.testing.assert_close( + tw, torch.full_like(tw, 1 / TOPK).bfloat16().float(), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize( + "num_tokens", + [ + 1, + moe.DECODE_MAX_TOKENS - 1, + moe.DECODE_MAX_TOKENS, + moe.DECODE_MAX_TOKENS + 1, + 64, + 600, + ], +) +@torch.inference_mode() +def test_moe_matches_reference(weights, flint_weights, num_tokens: int): + """Output follows the bf16 rounding rule (bitwise on >= 99% of elements) + and stays at the bf16 rounding level of the fp32 reference.""" + x, shared = _inputs(num_tokens) + out = _run(x, weights["router"], flint_weights, shared) + ids, tw = moe.route(x, weights["router"], TOPK, moe.make_route_counter("cuda")) + ref16 = _reference(x, weights, ids.long(), tw, shared, bf16_points=True) + ref32 = _reference(x, weights, ids.long(), tw, shared, bf16_points=False) + match = (out == ref16).float().mean().item() + assert match >= 0.99, match + assert _rel(out, ref32) < 1.25 * _rel(ref16, ref32) + tok_err = (out.float() - ref32).norm(dim=-1) / ref32.norm(dim=-1) + assert tok_err.max() < 0.01 + + +@torch.inference_mode() +def test_experts_module_loads_per_expert_tensors(weights, flint_weights): + """FlintNvfp4Experts assembles per-expert checkpoint tensors correctly and + repacks only once every distinct tensor is loaded.""" + layer = moe.FlintNvfp4Experts(E, H, INTER, TOPK).cuda() + + def ckpt(experts): + for e in experts: + for proj, rows in ( + ("gate_proj", slice(INTER)), + ("up_proj", slice(INTER, 2 * INTER)), + ): + yield f"{e}.{proj}.weight", weights["w13"][e, rows] + yield f"{e}.{proj}.weight_scale", weights["s13"][e, rows] + yield f"{e}.{proj}.weight_scale_2", weights["gs13"][e] + yield f"{e}.down_proj.weight", weights["w2"][e] + yield f"{e}.down_proj.weight_scale", weights["s2"][e] + yield f"{e}.down_proj.weight_scale_2", weights["gs2"][e] + + # Duplicates of expert 0 must not count towards completion. + layer.load_weights(ckpt([0] * E)) + assert layer.weights is None + layer.load_weights(ckpt(range(E))) + assert layer.weights is not None and not list(layer.parameters()) + x, shared = _inputs(4) + torch.testing.assert_close( + layer(x, weights["router"], shared), + _run(x, weights["router"], flint_weights, shared), + rtol=0, + atol=0, + ) + + +@pytest.mark.parametrize("num_tokens", [1, 64]) +@torch.inference_mode() +def test_shared_expert_on_aux_stream(weights, flint_weights, num_tokens: int): + """Passing the shared-expert module (run on the auxiliary stream, joined + by an event) gives the same output as passing its precomputed output.""" + layer = moe.FlintNvfp4Experts(E, H, INTER, TOPK).cuda() + layer.weights = flint_weights + for name in list(layer._parameters): + delattr(layer, name) + shared_mod = torch.nn.Linear(H, H, bias=False, device="cuda", dtype=torch.bfloat16) + x, _ = _inputs(num_tokens) + expected = layer(x, weights["router"], shared_mod(x)) + torch.testing.assert_close( + layer(x, weights["router"], shared_mod), expected, rtol=0, atol=0 + ) + + +@pytest.fixture(scope="module") +def bf16_weights(): + """BF16 experts as stored by vLLM's FlashInfer TRTLLM BF16 MoE backend, + plus the checkpoint-order originals.""" + import vllm.model_executor.layers.fused_moe # noqa: F401 (import order) + from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( + convert_moe_weights_to_flashinfer_trtllm_block_layout, + ) + + gen = torch.Generator(device="cuda").manual_seed(1) + w13 = (torch.randn(E, 2 * INTER, H, device="cuda", generator=gen) * 0.03).bfloat16() + w2 = (torch.randn(E, H, INTER, device="cuda", generator=gen) * 0.03).bfloat16() + ref = (w13.clone(), w2.clone()) + convert_moe_weights_to_flashinfer_trtllm_block_layout( + {}, w13, w2, is_gated_act_gemm=True + ) + return w13, w2, ref + + +def _trtllm_bf16(x, ids, tw, w13, w2): + import flashinfer + from flashinfer.fused_moe import WeightLayout + + from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import ( + view_as_block_major_k, + ) + + out = flashinfer.fused_moe.trtllm_bf16_routed_moe( + topk_ids=(ids, tw), + hidden_states=x, + gemm1_weights=view_as_block_major_k(w13), + gemm2_weights=view_as_block_major_k(w2), + num_experts=E, + top_k=TOPK, + n_group=None, + topk_group=None, + intermediate_size=INTER, + local_expert_offset=0, + local_num_experts=E, + routed_scaling_factor=None, + routing_method_type=1, + use_shuffled_weight=True, + weight_layout=WeightLayout.BlockMajorK, + do_finalize=True, + activation_type=3, # SwiGLU + ) + return out[0] if isinstance(out, list) else out + + +@pytest.mark.parametrize("num_tokens", [1, moe.BF16_DECODE_MAX_TOKENS]) +@torch.inference_mode() +def test_bf16_decode_matches_trtllm(weights, bf16_weights, num_tokens: int): + """The BF16 decode path reads the TRTLLM layout correctly and follows the + TRTLLM BF16 MoE rounding points: on identical routing it agrees with the + TRTLLM kernels used for larger batches up to GEMM accumulation order, and + with an fp32-GEMM reference of the same rule.""" + w13, w2, (w13_ref, w2_ref) = bf16_weights + maps = moe.trtllm_bf16_row_maps(H, INTER, "cuda") + x, shared = _inputs(num_tokens) + counter = moe.make_route_counter("cuda") + out = moe.flint_moe_bf16_decode( + x, weights["router"], w13, w2, maps, TOPK, counter, shared + ) + ids, tw = moe.route(x, weights["router"], TOPK, counter) + bf = torch.bfloat16 + trt = (_trtllm_bf16(x, ids, tw, w13, w2).float() + shared.float()).to(bf) + routed = torch.zeros(num_tokens, H, device="cuda") + for e in ids.unique().tolist(): + tok, slot = (ids == e).nonzero(as_tuple=True) + gu = x[tok].float() @ w13_ref[e].float().T + g, u = gu[:, :INTER], gu[:, INTER:] + h = (torch.nn.functional.silu(g) * u).to(bf).float() + y = (h @ w2_ref[e].float().T).to(bf).float() + routed.index_add_(0, tok.long(), tw[tok, slot, None] * y) + ref = (routed.to(bf).float() + shared.float()).to(bf) + for other in (trt, ref): + match = (out == other).float().mean().item() + assert match >= 0.97, match + assert _rel(out, other) < 1e-3, _rel(out, other) + + +def _shared_nvfp4(gen): + from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + swizzle_blockscale, + ) + + out = {} + for name, (rows, cols) in (("gu", (2 * 1024, H)), ("down", (H, 1024))): + w = torch.randint( + 0, 256, (rows, cols // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + exp = torch.randint(4, 10, (rows, cols // 16), device="cuda", generator=gen) + s = (exp << 3).to(torch.uint8).view(torch.float8_e4m3fn) + gs = torch.tensor([1.5e-3], device="cuda") + flat = swizzle_blockscale(s).contiguous().view(torch.uint8).view(-1) + out[name] = (w, s, gs, flat) + return out + + +def _ll_weights(weights, flint_weights, sh): + return moe.LLMoEWeights( + router=weights["router"], + experts=flint_weights, + shared_gu=sh["gu"][0].view(torch.int32), + shared_gu_scale=sh["gu"][3], + shared_gu_gs=sh["gu"][2], + shared_down=sh["down"][0].view(torch.int32), + shared_down_scale=sh["down"][3], + shared_down_gs=sh["down"][2], + ) + + +@pytest.mark.parametrize("num_tokens", [1, 2]) +@torch.inference_mode() +def test_ll_moe_block(weights, flint_weights, num_tokens): + """The fused decode MoE block equals Flint's fused_add_rms_norm followed by + the shared expert and the regular Flint decode MoE path.""" + from vllm.models.flint.ops import norm + + gen = torch.Generator(device="cuda").manual_seed(5) + bf = torch.bfloat16 + sh = _shared_nvfp4(gen) + lw = _ll_weights(weights, flint_weights, sh) + ln_w = (torch.rand(H, device="cuda", generator=gen) + 0.5).to(bf) + attn = (torch.randn(num_tokens, H, device="cuda", generator=gen) * 0.5).to(bf) + res = (torch.randn(num_tokens, H, device="cuda", generator=gen) * 2).to(bf) + counter = moe.make_route_counter("cuda") + out, res_out = moe.flint_moe_ll(attn, res, ln_w, 1e-6, lw, TOPK, counter) + x, r = norm.fused_add_rms_norm(attn.clone(), res.clone(), ln_w, 1e-6) + torch.testing.assert_close(res_out, r, rtol=0, atol=0) + gu = (x.float() @ dequantize_nvfp4(*sh["gu"][:3]).T).to(bf).float() + h = (torch.nn.functional.silu(gu[:, :1024]) * gu[:, 1024:]).to(bf) + shared = (h.float() @ dequantize_nvfp4(*sh["down"][:3]).T).to(bf) + ref = moe.flint_moe( + x, + weights["router"], + flint_weights, + TOPK, + moe.make_route_counter("cuda"), + shared, + ) + match = (out == ref).float().mean().item() + assert match >= 0.98, match + assert _rel(out, ref) < 1e-3 + + +@pytest.mark.parametrize("preset", [-(2**31), -128, 2**31 - 128]) +@torch.inference_mode() +def test_router_counters_wrap(weights, flint_weights, preset): + """The never-reset router counters elect the token's last program + correctly across signed 32-bit wraparound (route() and the fused block).""" + x, shared = _inputs(2, seed=3) + counter = moe.make_route_counter("cuda") + ref_ids, ref_tw = moe.route(x, weights["router"], TOPK, counter) + counter.fill_(preset) + for _ in range(3): + ids, tw = moe.route(x, weights["router"], TOPK, counter) + torch.testing.assert_close(ids, ref_ids, rtol=0, atol=0) + torch.testing.assert_close(tw, ref_tw, rtol=0, atol=0) + + gen = torch.Generator(device="cuda").manual_seed(5) + lw = _ll_weights(weights, flint_weights, _shared_nvfp4(gen)) + ln_w = torch.ones(H, device="cuda", dtype=torch.bfloat16) + counter = moe.make_route_counter("cuda") + ref, _ = moe.flint_moe_ll(x, shared, ln_w, 1e-6, lw, TOPK, counter) + counter.fill_(preset) + for _ in range(3): + out, _ = moe.flint_moe_ll(x, shared, ln_w, 1e-6, lw, TOPK, counter) + torch.testing.assert_close(out, ref, rtol=0, atol=0) diff --git a/tests/models/flint/test_flint_ops.py b/tests/models/flint/test_flint_ops.py new file mode 100644 index 00000000000..ad441240df2 --- /dev/null +++ b/tests/models/flint/test_flint_ops.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint's dense-linear and elementwise ops against high-precision references: +the small-M NVFP4 / BF16 GEMVs (``nvfp4_gemv``, incl. the fused SiLU-and-mul +epilogue), the large-M dequantize + cuBLAS path and global-scale check +(``linear``), and the Triton norms and SiLU-and-mul (``norm``).""" + +import pytest +import torch + +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + swizzle_blockscale, +) +from vllm.models.flint.gguf_nvfp4 import dequantize_nvfp4 +from vllm.models.flint.ops import nvfp4_gemv +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="SM100 only", +) + + +def _cases(table): + """(N, K, M) for each tuned shape: M within the first token tile, and a + partial and a full largest tile.""" + return [ + (n, k, m) + for (n, k), cfgs in table.items() + for m in sorted({1, 5, 8 * max(cfgs) - 3, 8 * max(cfgs)}) + ] + + +def _nvfp4(n: int, k: int, seed: int = 0): + gen = torch.Generator(device="cuda").manual_seed(seed) + w = torch.randint( + 0, 256, (n, k // 2), dtype=torch.uint8, device="cuda", generator=gen + ) + # Every finite non-negative E4M3 code, subnormals and zero included. + s = torch.randint( + 0, 127, (n, k // 16), dtype=torch.uint8, device="cuda", generator=gen + ).view(torch.float8_e4m3fn) + gs = torch.tensor(1.7e-3, device="cuda") + return w, s, gs + + +def _reference(x, w, s, gs): + return x.double() @ dequantize_nvfp4(w, s, gs).double().t() + + +def _silu_and_mul_ref(gate_up): + g, u = gate_up.float().chunk(2, dim=-1) + return (g / (1 + torch.exp(-g)) * u).bfloat16() + + +def _check(out, ref): + # FP32 accumulation over K <= 6144 terms vs. a float64 reference: far + # below one BF16 ulp, so the only visible error is the output rounding. + err = (out.double() - ref).abs() + tol = 2.0**-8 * ref.abs() + 1e-3 * ref.abs().max() + assert (err <= tol).all(), f"max err {err.max().item()}" + + +@pytest.mark.parametrize("n,k,m", _cases(nvfp4_gemv._TUNED)) +def test_gemv_matches_dequantized_reference(n, k, m): + w, s, gs = _nvfp4(n, k) + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + out = nvfp4_gemv.nvfp4_gemv(x, w, swizzle_blockscale(s), gs) + _check(out, _reference(x, w, s, gs)) + + +def test_gemv_dequantizes_every_code_exactly(): + """A one-hot activation reads back single dequantized weights bit-exactly: + all 16 E2M1 codes times all finite E4M3 scales, incl. subnormals.""" + n, k = 2048, 4096 + w = torch.arange(256, dtype=torch.uint8, device="cuda").repeat(n, k // 512) + w = w[torch.randperm(n, device="cuda")] + s = (torch.arange(n * k // 16, device="cuda") % 127).to(torch.uint8) + s = s.reshape(n, k // 16).view(torch.float8_e4m3fn) + gs = torch.tensor(1.7e-3, device="cuda") + m = 16 + x = torch.zeros(m, k, dtype=torch.bfloat16, device="cuda") + x[torch.arange(m), torch.randint(0, k, (m,))] = 1.0 + out = nvfp4_gemv.nvfp4_gemv(x, w, swizzle_blockscale(s), gs) + ref = (x.float() @ dequantize_nvfp4(w, s, gs).t()).bfloat16() + torch.testing.assert_close(out, ref, rtol=0, atol=0) + + +def test_gemv_strided_input(): + n, k = 2048, 4096 + w, s, gs = _nvfp4(n, k) + base = torch.randn(4, k + 256, dtype=torch.bfloat16, device="cuda") + x = base[:, 128 : 128 + k] + out = nvfp4_gemv.nvfp4_gemv(x, w, swizzle_blockscale(s), gs) + _check(out, _reference(x, w, s, gs)) + + +@pytest.mark.parametrize("n,k,m", _cases(nvfp4_gemv._TUNED_SILU_MUL)) +def test_gemv_silu_mul_matches_unfused(n, k, m): + """The fused gate_up + SiLU-and-mul epilogue equals SiLU-and-mul applied to + the BF16-rounded GEMM output.""" + w, s, _ = _nvfp4(n, k) + gs = torch.tensor(1e-5, device="cuda") # gate/up values of O(1-10) + gen = torch.Generator(device="cuda").manual_seed(m) + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda", generator=gen) + ref = _silu_and_mul_ref(_reference(x, w, s, gs).bfloat16()) + out = nvfp4_gemv.nvfp4_gemv(x, w, swizzle_blockscale(s), gs, silu_mul=True) + # FP32 vs float64 sums can round a gate/up value one BF16 ulp apart, which + # SiLU amplifies for negative gates (~3% at -8). + err = (out.float() - ref.float()).abs() + close = err <= 2**-7 * ref.float().abs() + 1e-6 + assert close.float().mean() > 0.999 + assert (err <= 0.1 * ref.float().abs() + 1e-5).all() + + +def test_merged_global_scale_mismatch_raises(): + from vllm.models.flint.ops.linear import check_merged_global_scales + + p = "model.layers.0.self_attn." + ok = [(f"{p}{n}_proj.weight_scale_2", torch.tensor(0.5)) for n in "qkv"] + assert len(list(check_merged_global_scales(ok))) == 3 + bad = ok[:2] + [(p + "v_proj.weight_scale_2", torch.tensor(0.25))] + with pytest.raises(ValueError): + list(check_merged_global_scales(bad)) + + +@pytest.mark.parametrize("n,k,m", _cases(nvfp4_gemv._TUNED_BF16)) +def test_bf16_gemv_matches_reference(n, k, m): + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") + out = nvfp4_gemv.bf16_gemv(x, w) + _check(out, x.double() @ w.double().t()) + + +def test_bf16_gemv_silu_mul(): + n, k = 2048, 2048 + gen = torch.Generator(device="cuda").manual_seed(0) + w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda", generator=gen) * 0.05 + x = torch.randn(4, k, dtype=torch.bfloat16, device="cuda", generator=gen) + ref = _silu_and_mul_ref((x.double() @ w.double().t()).bfloat16()) + out = nvfp4_gemv.bf16_gemv(x, w, silu_mul=True) + err = (out.float() - ref.float()).abs() + assert (err <= 2**-7 * ref.float().abs() + 1e-6).float().mean() > 0.999 + + +def test_large_m_dequant_gemm_matches_reference(): + from types import SimpleNamespace + + from vllm.models.flint.ops.linear import _dequant_gemm + + n, k = 2048, 4096 + w, s, gs = _nvfp4(n, k) + layer = SimpleNamespace( + weight=w, + flint_weight_scale=swizzle_blockscale(s), + flint_global_scale=float(gs), + ) + x = torch.randn(256, k, dtype=torch.bfloat16, device="cuda") + _check(_dequant_gemm(layer, x), _reference(x, w, s, gs)) + + +def _rms_norm_ref(s, w, eps): + rs = torch.rsqrt(s.pow(2).mean(-1, keepdim=True) + eps) + return (s * rs * w.float()).bfloat16() + + +@pytest.mark.parametrize("m", [1, 7, 8192]) +@pytest.mark.parametrize("has_residual", [False, True]) +@torch.inference_mode() +def test_rms_norm_matches_reference(m, has_residual): + """Flint's (fused residual-add) RMSNorm: fp32 math on the unrounded + x + residual, one bf16 rounding of each output. rsqrt.approx and the + reduction order allow rare 1-ulp differences from the fp32 reference.""" + from vllm.models.flint.ops import norm + + gen = torch.Generator(device="cuda").manual_seed(m) + h = 2048 + x = (torch.randn(m, h, device="cuda", generator=gen) * 2).bfloat16() + w = torch.empty(h, device="cuda").normal_(1, 0.2, generator=gen).bfloat16() + if has_residual: + r = (torch.randn(m, h, device="cuda", generator=gen) * 4).bfloat16() + s = x.float() + r.float() + out, res_out = norm.fused_add_rms_norm(x.clone(), r.clone(), w, 1e-6) + torch.testing.assert_close(res_out, s.bfloat16(), rtol=0, atol=0) + else: + out = norm.rms_norm(x, w, 1e-6) + s = x.float() + ref = _rms_norm_ref(s, w, 1e-6) + assert (out == ref).float().mean() > 0.999 + torch.testing.assert_close(out, ref, rtol=2**-7, atol=0) + + +@pytest.mark.parametrize("m", [1, 7, 8192]) +@torch.inference_mode() +def test_silu_and_mul_matches_reference(m): + """bf16(silu(g) * u) in fp32 with one rounding.""" + from vllm.models.flint.ops import norm + + gen = torch.Generator(device="cuda").manual_seed(m) + x = (torch.randn(m, 2 * 6144, device="cuda", generator=gen) * 3).bfloat16() + out = norm.silu_and_mul(x) + ref = _silu_and_mul_ref(x) + assert (out == ref).float().mean() > 0.999 + torch.testing.assert_close(out, ref, rtol=2**-7, atol=1e-6) diff --git a/tests/models/flint/test_gguf_nvfp4.py b/tests/models/flint/test_gguf_nvfp4.py new file mode 100644 index 00000000000..09876a5b0e7 --- /dev/null +++ b/tests/models/flint/test_gguf_nvfp4.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint's NVFP4 GGUF loader must repack llama.cpp NVFP4 blocks into the +ModelOpt layout without changing any weight value.""" + +import numpy as np +import pytest +import torch + +from vllm.models.flint import gguf_nvfp4 +from vllm.models.flint.gguf_nvfp4 import ( + _repack_nvfp4, + dequantize_nvfp4, + find_gguf_file, + gguf_nvfp4_weights_iterator, +) + +gguf = pytest.importorskip("gguf", minversion="0.19.0") + +NVFP4 = gguf.GGMLQuantizationType.NVFP4 + + +def _random_blocks(rows: int, cols: int, seed: int = 0) -> np.ndarray: + # Random bytes cover every E2M1 code and every representable scale byte, + # including the sign bit and the 0x7F NaN pattern that ggml decodes as 0. + rng = np.random.default_rng(seed) + blocks = rng.integers(0, 256, (rows, cols // 64, 36), dtype=np.uint8) + blocks[..., :4] = np.where(blocks[..., :4] == 0xFF, 0x7F, blocks[..., :4]) + return blocks.reshape(rows, -1) + + +def test_repack_matches_gguf_dequantize(): + rows, cols = 32, 512 + blocks = _random_blocks(rows, cols) + weight, weight_scale = _repack_nvfp4(torch.from_numpy(blocks), cols) + assert weight.shape == (rows, cols // 2) + assert weight_scale.shape == (rows, cols // 16) + ours = dequantize_nvfp4(weight, weight_scale) + ref = gguf.quants.dequantize(blocks, NVFP4).reshape(rows, cols) + torch.testing.assert_close(ours, torch.from_numpy(ref), atol=0, rtol=0) + + +def test_unrepresentable_scale_raises(): + blocks = _random_blocks(1, 64) + blocks[0, 0] = 0xFF + with pytest.raises(ValueError): + _repack_nvfp4(torch.from_numpy(blocks), 64) + + +def test_iterator_maps_names_and_scales(tmp_path): + hidden, inter, vocab, experts = 128, 64, 256, 4 + path = tmp_path / "tiny.gguf" + writer = gguf.GGUFWriter(str(path), "flint") + + def add_nvfp4(name, rows, cols, num, seed): + blocks = _random_blocks(num * rows, cols, seed).reshape(num, rows, -1) + data = blocks[0] if num == 1 else blocks + writer.add_tensor(name + ".weight", data, raw_dtype=NVFP4) + scale = np.full((num,), 0.5, dtype=np.float32) * np.arange(1, num + 1) + writer.add_tensor(name + ".scale", scale) + return data + + q = add_nvfp4("blk.2.attn_q", hidden, hidden, 1, 1) + add_nvfp4("blk.2.ffn_up_exps", inter, hidden, experts, 2) + head = add_nvfp4("output", vocab, hidden, 1, 3) + writer.add_tensor("blk.2.attn_sinks.weight", np.zeros(4, dtype=np.float32)) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + tensors = dict( + gguf_nvfp4_weights_iterator(str(path), "cpu", frozenset({"lm_head"})) + ) + + q_w = tensors["model.layers.2.self_attn.q_proj.weight"] + q_s = tensors["model.layers.2.self_attn.q_proj.weight_scale"] + assert q_w.dtype == torch.uint8 and q_s.dtype == torch.float8_e4m3fn + ref = gguf.quants.dequantize(q, NVFP4).reshape(hidden, hidden) * 0.5 + torch.testing.assert_close( + dequantize_nvfp4( + q_w, q_s, tensors["model.layers.2.self_attn.q_proj.weight_scale_2"] + ), + torch.from_numpy(ref), + atol=0, + rtol=0, + ) + # Stacked experts are split per expert with their own global scale. + for e in range(experts): + gs = tensors[f"model.layers.2.mlp.experts.{e}.up_proj.weight_scale_2"] + assert float(gs) == 0.5 * (e + 1) + # The LM head is dequantized exactly to bf16 (global scale kept apart). + head_w = tensors["lm_head.weight"] + assert head_w.dtype == torch.bfloat16 + head_ref = gguf.quants.dequantize(head, NVFP4).reshape(vocab, hidden) + torch.testing.assert_close( + head_w.float(), torch.from_numpy(head_ref), atol=0, rtol=0 + ) + assert float(tensors["lm_head.weight_scale_2"]) == 0.5 + assert "model.layers.2.self_attn.softmax_offset" in tensors + + +def test_find_gguf_file_resolves_hub_repo_ids(tmp_path, monkeypatch): + """`vllm serve /` must detect a GGUF-only repo, not just local paths.""" + (tmp_path / "flint.gguf").touch() + monkeypatch.setattr( + gguf_nvfp4, + "any_pattern_in_repo_files", + lambda repo, patterns, revision: repo == "org/gguf", + ) + monkeypatch.setattr( + gguf_nvfp4, "download_weights_from_hf", lambda *args: str(tmp_path) + ) + assert find_gguf_file("org/gguf") == str(tmp_path / "flint.gguf") + assert find_gguf_file("org/safetensors") is None diff --git a/tests/models/flint/test_qk_norm_rope.py b/tests/models/flint/test_qk_norm_rope.py new file mode 100644 index 00000000000..c3fd19b0bd6 --- /dev/null +++ b/tests/models/flint/test_qk_norm_rope.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The fused q/k norm + RoPE + cache-write kernel must reproduce the unfused +vLLM ops (rms_norm, rotary_embedding, reshape_and_cache_flash) it replaces.""" + +import pytest +import torch + +from vllm import _custom_ops as ops +from vllm.models.flint.ops.qk_norm_rope import qk_norm_rope_cache +from vllm.platforms import current_platform + +HQ, HKV, D = 32, 4, 128 +BLOCK_SIZE, NUM_BLOCKS = 16, 64 + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("num_tokens", [1, 7, 256, 1000]) +@pytest.mark.parametrize("apply_rope", [True, False]) +@torch.inference_mode() +def test_matches_unfused_ops(num_tokens: int, apply_rope: bool): + torch.manual_seed(0) + dev = "cuda" + qkv = torch.randn(num_tokens, (HQ + 2 * HKV) * D, device=dev).bfloat16() * 3 + positions = torch.randint(0, 100_000, (num_tokens,), device=dev) + q_w = torch.empty(D, device=dev).normal_(1.0, 0.2).bfloat16() + k_w = torch.empty(D, device=dev).normal_(1.0, 0.2).bfloat16() + inv_freq = 1.0 / (10000.0 ** (torch.arange(0, D, 2, device=dev).float() / D)) + freqs = torch.outer(torch.arange(262144, device=dev).float(), inv_freq) + cos_sin = torch.cat([freqs.cos(), freqs.sin()], dim=-1) + # Padded tokens (beyond the slot mapping) and slot -1 must not be written. + slots = torch.randperm(NUM_BLOCKS * BLOCK_SIZE, device=dev)[:num_tokens] + slot_mapping = slots[: max(1, num_tokens - 1)].clone() + slot_mapping[0] = -1 + + q, k, v = qkv.split([HQ * D, HKV * D, HKV * D], dim=-1) + ref_q = torch.empty(num_tokens, HQ, D, device=dev, dtype=torch.bfloat16) + ref_k = torch.empty(num_tokens, HKV, D, device=dev, dtype=torch.bfloat16) + ops.rms_norm(ref_q, q.reshape(-1, HQ, D), q_w, 1e-6) + ops.rms_norm(ref_k, k.reshape(-1, HKV, D), k_w, 1e-6) + ref_q = ref_q.reshape(q.shape) + ref_k = ref_k.reshape(k.shape) + if apply_rope: + # fp32 cos/sin with bf16 q/k, as in the reference implementation. + ops.rotary_embedding(positions, ref_q, ref_k, D, cos_sin, True) + ref_cache = torch.zeros(NUM_BLOCKS, 2, BLOCK_SIZE, HKV, D, device=dev).bfloat16() + torch.ops._C_cache_ops.reshape_and_cache_flash( + ref_k.reshape(-1, HKV, D), + v.reshape(-1, HKV, D), + ref_cache[:, 0], + ref_cache[:, 1], + slot_mapping, + "auto", + torch.ones((), device=dev), + torch.ones((), device=dev), + ) + + cache = torch.zeros_like(ref_cache) + out_qkv = qkv.clone() + out_q = qk_norm_rope_cache( + out_qkv, + positions, + q_w, + k_w, + 1e-6, + cos_sin if apply_rope else None, + cache[:, 0], + cache[:, 1], + slot_mapping, + HQ, + HKV, + D, + ) + out_k = out_qkv[:, HQ * D : (HQ + HKV) * D] + # Same ops and roundings; only the RMS reduction order differs. + torch.testing.assert_close(out_q, ref_q, atol=4e-3, rtol=1.6e-2) + torch.testing.assert_close(out_k, ref_k, atol=4e-3, rtol=1.6e-2) + assert (out_q != ref_q).float().mean() < 0.01 + assert (out_k != ref_k).float().mean() < 0.01 + torch.testing.assert_close(cache, ref_cache, atol=4e-3, rtol=1.6e-2) + assert torch.equal(out_qkv[:, (HQ + HKV) * D :], v) diff --git a/tests/models/flint/test_qkv_fused.py b/tests/models/flint/test_qkv_fused.py new file mode 100644 index 00000000000..9bc04887c78 --- /dev/null +++ b/tests/models/flint/test_qkv_fused.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The two-kernel fused decode qkv (residual-add RMSNorm, then qkv GEMV with +per-head q/k norm, RoPE and KV-cache write) must reproduce the kernels it +replaces, Flint's ``fused_add_rms_norm`` / ``rms_norm`` -> Flint GEMV -> +``qk_norm_rope_cache``, up to reduction order (rare 1-ulp differences).""" + +import pytest +import torch + +from vllm.platforms import current_platform + +if not current_platform.is_cuda() or not current_platform.is_device_capability_family( + 100 +): + pytest.skip("Flint fused qkv targets SM100", allow_module_level=True) + +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( # noqa: E402 + swizzle_blockscale, +) +from vllm.models.flint.ops import norm, nvfp4_gemv # noqa: E402 +from vllm.models.flint.ops.qk_norm_rope import qk_norm_rope_cache # noqa: E402 +from vllm.models.flint.ops.qkv_fused import flint_qkv_fused # noqa: E402 + +HQ, HKV, D, K = 32, 4, 128, 2048 +N = (HQ + 2 * HKV) * D +NUM_BLOCKS, BLOCK_SIZE = 64, 16 + + +def _weights(bf16w: bool, g: torch.Generator): + if bf16w: + w = (torch.randn(N, K, device="cuda", generator=g) * 0.02).bfloat16() + return w, None, None + w = torch.randint(0, 256, (N, K // 2), device="cuda", generator=g) + scales = torch.rand(N, K // 16, device="cuda", generator=g) * 2 + 0.25 + scales = swizzle_blockscale(scales.to(torch.float8_e4m3fn)) + return w.to(torch.uint8), scales, torch.tensor([0.0037], device="cuda") + + +@pytest.mark.parametrize("m", [1, 2, 5, 8]) +@pytest.mark.parametrize("bf16w", [False, True], ids=["nvfp4", "bf16"]) +@pytest.mark.parametrize( + "has_residual,rope", + [(True, True), (True, False), (False, True), (False, False)], + ids=["res-rope", "res-norope", "first-layer", "first-norope"], +) +@pytest.mark.parametrize("short_slots", [False, True], ids=["pad", "short"]) +@torch.inference_mode() +def test_matches_unfused_kernels( + m: int, bf16w: bool, has_residual: bool, rope: bool, short_slots: bool +): + g = torch.Generator(device="cuda").manual_seed(m) + w, ws, gs = _weights(bf16w, g) + x = (torch.randn(m, K, device="cuda", generator=g) * 3).bfloat16() + res = None + if has_residual: + res = (torch.randn(m, K, device="cuda", generator=g) * 3).bfloat16() + norm_w = torch.empty(K, device="cuda").normal_(1, 0.2, generator=g).bfloat16() + q_w = torch.empty(D, device="cuda").normal_(1, 0.2, generator=g).bfloat16() + k_w = torch.empty(D, device="cuda").normal_(1, 0.2, generator=g).bfloat16() + cos_sin = None + if rope: + inv = 1.0 / (10000.0 ** (torch.arange(0, D, 2, device="cuda").float() / D)) + freqs = torch.outer(torch.arange(262144, device="cuda").float(), inv) + # FP32 cos | sin table, as get_rope(..., dtype=torch.float32) builds. + cos_sin = torch.cat([freqs.cos(), freqs.sin()], dim=-1) + # Up to the end of the RoPE table. + positions = torch.randint(262144 - 64, 262144, (m,), device="cuda", generator=g) + slots = torch.randperm(NUM_BLOCKS * BLOCK_SIZE, device="cuda", generator=g)[:m] + # Padding tokens are skipped: a -1 slot, or beyond a shorter slot mapping. + if m > 1 and short_slots: + slots = slots[:-1] + elif m > 1: + slots[-1] = -1 + kv = torch.zeros( + NUM_BLOCKS, HKV, BLOCK_SIZE, 2 * D, device="cuda", dtype=torch.bfloat16 + ) + ref_kv = kv.clone() + + if res is not None: + xr, ref_res = norm.fused_add_rms_norm(x.clone(), res.clone(), norm_w, 1e-6) + else: + xr = norm.rms_norm(x, norm_w, 1e-6) + if bf16w: + qkv = (xr.float() @ w.float().t()).bfloat16() + else: + qkv = nvfp4_gemv.nvfp4_gemv(xr, w, ws, gs) + ref_k, ref_v = ref_kv.transpose(1, 2).split(D, dim=-1) + ref_q = qk_norm_rope_cache( + qkv, positions, q_w, k_w, 1e-6, cos_sin, ref_k, ref_v, slots, HQ, HKV, D + ) + + k_cache, v_cache = kv.transpose(1, 2).split(D, dim=-1) + q, res_out = flint_qkv_fused( + x, res, norm_w, 1e-6, w, ws, gs, q_w, k_w, 1e-6, cos_sin, positions, + k_cache, v_cache, slots, HQ, HKV, + ) # fmt: skip + + if has_residual: + torch.testing.assert_close(res_out, ref_res, atol=0, rtol=0) + else: + assert res_out is x + # Same rounding points, but the norms (and the BF16 GEMV's K split) sum in + # a different fp32 order: a 1-ulp difference in a normalized activation + # shifts every qkv output of its token slightly. + assert (q == ref_q).float().mean() > 0.9 + assert (kv == ref_kv).float().mean() > 0.9 + torch.testing.assert_close(q, ref_q, atol=1e-2, rtol=2e-2) + torch.testing.assert_close(kv, ref_kv, atol=1e-2, rtol=2e-2) diff --git a/tests/models/registry.py b/tests/models/registry.py index 6c4d548800b..342ec48c70d 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -278,6 +278,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { ), "FalconForCausalLM": _HfExamplesInfo("tiiuae/falcon-7b"), "FalconH1ForCausalLM": _HfExamplesInfo("tiiuae/Falcon-H1-0.5B-Base"), + "FlintForCausalLM": _HfExamplesInfo("RecursiveSI/flint", is_available_online=False), "FlexOlmoForCausalLM": _HfExamplesInfo("allenai/Flex-reddit-2x7B-1T"), "GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"), "Gemma2ForCausalLM": _HfExamplesInfo( diff --git a/tests/transformers_utils/test_speculators_override.py b/tests/transformers_utils/test_speculators_override.py index 1495d5a87c4..b3de532d0fa 100644 --- a/tests/transformers_utils/test_speculators_override.py +++ b/tests/transformers_utils/test_speculators_override.py @@ -235,3 +235,12 @@ def test_speculators_rejects_local_verifier_before_model_loading( with pytest.raises(ValueError, match="Speculators verifier must"): engine_args_cls(model=str(draft)).create_engine_config() + + +def test_local_file_model_skips_speculators_detection(tmp_path): + """A single-file checkpoint (e.g. a .gguf) is not parsed as a config.""" + gguf = tmp_path / "model.gguf" + gguf.write_bytes(b"GGUF") + assert maybe_override_with_speculators( + model=str(gguf), tokenizer="org/tok", trust_remote_code=False + ) == (str(gguf), "org/tok", None) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 89f39a40ffa..cb92b17306b 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -106,6 +106,7 @@ _TEXT_GENERATION_MODELS = { "FalconForCausalLM": ("falcon", "FalconForCausalLM"), "FalconMambaForCausalLM": ("mamba", "MambaForCausalLM"), "FalconH1ForCausalLM": ("falcon_h1", "FalconH1ForCausalLM"), + "FlintForCausalLM": ("vllm.models.flint", "FlintForCausalLM"), "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"), "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"), "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"), diff --git a/vllm/models/flint/__init__.py b/vllm/models/flint/__init__.py new file mode 100644 index 00000000000..8f80a50a830 --- /dev/null +++ b/vllm/models/flint/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint model entry point.""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .model import FlintForCausalLM + +__all__ = ["FlintForCausalLM"] + + +def __getattr__(name: str): + # Lazy so that resolving FlintConfig does not import CUDA kernels. + if name in __all__: + from . import model + + return getattr(model, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm/models/flint/configs.py b/vllm/models/flint/configs.py new file mode 100644 index 00000000000..62a69e325cf --- /dev/null +++ b/vllm/models/flint/configs.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint model config.""" + +from typing import Any + +from transformers.configuration_utils import PreTrainedConfig + + +class FlintConfig(PreTrainedConfig): + """Decoder-only MoE transformer with hybrid sliding/full attention. + + Sliding layers use RoPE; full-attention layers listed in + ``no_rope_layer_indices`` use no positional encoding. Every attention head + has a learned sink logit. ``sliding_window`` counts the current token. + """ + + model_type = "flint" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 150144, + hidden_size: int = 2048, + intermediate_size: int = 6144, + num_hidden_layers: int = 49, + num_attention_heads: int = 32, + num_key_value_heads: int = 4, + head_dim: int = 128, + hidden_act: str = "silu", + max_position_embeddings: int = 262144, + initializer_range: float = 0.02, + rms_norm_eps: float = 1e-6, + use_cache: bool = True, + tie_word_embeddings: bool = False, + rope_theta: float = 10000.0, + rope_parameters: dict[str, Any] | None = None, + attention_bias: bool = False, + attention_dropout: float = 0.0, + use_sliding_window: bool = True, + sliding_window: int | None = 513, + layer_types: list[str] | None = None, + no_rope_layer_indices: list[int] | None = None, + moe_intermediate_size: int = 512, + shared_expert_intermediate_size: int = 1024, + num_experts_per_tok: int = 6, + num_experts: int = 256, + norm_topk_prob: bool = True, + router_score_function: str = "sigmoid", + mlp_only_layers: list[int] | None = None, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_parameters = rope_parameters or { + "rope_type": "default", + "rope_theta": rope_theta, + } + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if use_sliding_window else None + self.layer_types = layer_types or ["full_attention"] * num_hidden_layers + if len(self.layer_types) != num_hidden_layers: + raise ValueError( + f"layer_types has {len(self.layer_types)} entries but " + f"num_hidden_layers is {num_hidden_layers}" + ) + self.no_rope_layer_indices = no_rope_layer_indices or [] + self.moe_intermediate_size = moe_intermediate_size + self.shared_expert_intermediate_size = shared_expert_intermediate_size + self.num_experts_per_tok = num_experts_per_tok + self.num_experts = num_experts + self.norm_topk_prob = norm_topk_prob + self.router_score_function = router_score_function + self.mlp_only_layers = mlp_only_layers or [] + # The reference implementation computes the LM head in fp32. + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + + +__all__ = ["FlintConfig"] diff --git a/vllm/models/flint/gguf_nvfp4.py b/vllm/models/flint/gguf_nvfp4.py new file mode 100644 index 00000000000..ca0a0e46939 --- /dev/null +++ b/vllm/models/flint/gguf_nvfp4.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Load Flint's NVFP4 GGUF checkpoint as ModelOpt-layout NVFP4 tensors. + +A GGUF NVFP4 super-block holds 64 values as four 16-value sub-blocks, each with +a UE4M3 scale byte and 8 bytes of E2M1 codes (low nibbles = values 0..7, high +nibbles = values 8..15). The value is ``e4m3(scale) * e2m1(code) * tensor_scale`` +where ``tensor_scale`` is the companion ``.scale`` F32 tensor (one per +expert for stacked expert tensors). ModelOpt packs consecutive values into one +byte (low nibble first) and stores the same three factors as ``weight``, +``weight_scale`` and ``weight_scale_2``, so conversion is a lossless repack. +""" + +import glob +import os +import warnings +from collections.abc import Iterator + +import huggingface_hub +import numpy as np +import regex as re +import torch + +from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf +from vllm.transformers_utils.repo_utils import any_pattern_in_repo_files + +_GGUF_NVFP4 = 40 +_QK = 64 +_BLOCK_BYTES = 36 + +_LAYER_NAMES = { + "attn_norm.weight": "input_layernorm.weight", + "ffn_norm.weight": "post_attention_layernorm.weight", + "attn_q_norm.weight": "self_attn.q_norm.weight", + "attn_k_norm.weight": "self_attn.k_norm.weight", + "attn_sinks.weight": "self_attn.softmax_offset", + "attn_q": "self_attn.q_proj", + "attn_k": "self_attn.k_proj", + "attn_v": "self_attn.v_proj", + "attn_output": "self_attn.o_proj", + "ffn_gate": "mlp.gate_proj", + "ffn_up": "mlp.up_proj", + "ffn_down": "mlp.down_proj", + "ffn_gate_shexp": "mlp.shared_expert.gate_proj", + "ffn_up_shexp": "mlp.shared_expert.up_proj", + "ffn_down_shexp": "mlp.shared_expert.down_proj", + "ffn_gate_inp.weight": "mlp.gate.weight", +} +_EXPERT_NAMES = { + "ffn_gate_exps": "gate_proj", + "ffn_up_exps": "up_proj", + "ffn_down_exps": "down_proj", +} +_GLOBAL_NAMES = { + "token_embd.weight": "model.embed_tokens.weight", + "output_norm.weight": "model.norm.weight", + "output": "lm_head", +} + + +def find_gguf_file( + model_path: str, revision: str | None = None, download_dir: str | None = None +) -> str | None: + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + return model_path + if not os.path.isdir(model_path): + # Offline, the repo cannot be listed; look in the cached snapshot instead. + offline = huggingface_hub.constants.HF_HUB_OFFLINE + if not offline and not any_pattern_in_repo_files( + model_path, ["*.gguf"], revision=revision + ): + return None + try: + model_path = download_weights_from_hf( + model_path, download_dir, ["*.gguf"], revision + ) + except huggingface_hub.errors.LocalEntryNotFoundError: + return None + files = sorted(glob.glob(os.path.join(model_path, "*.gguf"))) + if len(files) > 1: + raise ValueError(f"Expected one GGUF file in {model_path}, found {files}") + return files[0] if files else None + + +def _repack_nvfp4( + blocks: torch.Tensor, n_cols: int +) -> tuple[torch.Tensor, torch.Tensor]: + """[..., n_cols / 64 * 36] GGUF bytes -> (packed [..., n_cols / 2] uint8, + block scales [..., n_cols / 16] float8_e4m3fn).""" + lead = blocks.shape[:-1] + blocks = blocks.reshape(*lead, n_cols // _QK, _BLOCK_BYTES) + # ggml decodes the scales as unsigned E4M3 and the NaN pattern as zero; + # 0xFF would decode to 480, which E4M3 cannot represent. + scales = blocks[..., :4] + if bool((scales == 0xFF).any()): + raise ValueError("NVFP4 GGUF block scale 0xFF is not supported") + scales = torch.where(scales == 0x7F, torch.zeros_like(scales), scales & 0x7F) + qs = blocks[..., 4:].reshape(*lead, n_cols // _QK, 4, 8) + codes = torch.cat([qs & 0x0F, qs >> 4], dim=-1) + packed = codes[..., 0::2] | (codes[..., 1::2] << 4) + return ( + packed.reshape(*lead, n_cols // 2).contiguous(), + scales.reshape(*lead, n_cols // 16).contiguous().view(torch.float8_e4m3fn), + ) + + +_E2M1_POS = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E2M1 = _E2M1_POS + tuple(-v for v in _E2M1_POS) + + +def dequantize_nvfp4( + weight: torch.Tensor, + weight_scale: torch.Tensor, + global_scale: torch.Tensor | None = None, + out_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ModelOpt-layout NVFP4 [N, K / 2] -> [N, K]. Without ``global_scale`` the + result is ``e2m1 * e4m3``, which is exact in bf16.""" + chunk_rows = 8192 + lut = torch.tensor(_E2M1, dtype=torch.float32, device=weight.device) + out = torch.empty( + (*weight.shape[:-1], weight.shape[-1] * 2), + dtype=out_dtype, + device=weight.device, + ) + for start in range(0, weight.shape[0], chunk_rows): + w = weight[start : start + chunk_rows] + codes = torch.stack([w & 0x0F, w >> 4], dim=-1).int() + values = lut[codes].reshape(*w.shape[:-1], -1, 16) + values = values * weight_scale[start : start + chunk_rows].float().unsqueeze(-1) + values = values.flatten(-2) + if global_scale is not None: + values = values * global_scale.float().to(weight.device) + out[start : start + chunk_rows] = values + return out + + +def _as_tensor(data: np.ndarray) -> torch.Tensor: + # The reader memory-maps the file read-only; tensors are only read. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + return torch.from_numpy(np.asarray(data)) + + +def _map_name(name: str) -> tuple[str, bool]: + """GGUF tensor base name -> (HF base name, is_stacked_experts).""" + if name in _GLOBAL_NAMES: + return _GLOBAL_NAMES[name], False + m = re.fullmatch(r"blk\.(\d+)\.(.+)", name) + if m is None: + raise ValueError(f"Unexpected tensor in Flint GGUF: {name}") + layer, rest = m.groups() + prefix = f"model.layers.{layer}." + if rest in _EXPERT_NAMES: + return prefix + "mlp.experts.{}." + _EXPERT_NAMES[rest], True + if rest not in _LAYER_NAMES: + raise ValueError(f"Unexpected tensor in Flint GGUF: {name}") + return prefix + _LAYER_NAMES[rest], False + + +def gguf_nvfp4_weights_iterator( + path: str, + device: torch.device | str = "cuda", + dequantize: frozenset[str] = frozenset(), +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield HF-style (name, tensor) pairs; NVFP4 tensors become ``.weight`` + (packed uint8), ``.weight_scale`` (E4M3) and ``.weight_scale_2`` (F32), + except HF base names in ``dequantize``, which become a BF16 ``.weight`` + holding ``e2m1 * e4m3`` (exact) plus the F32 ``.weight_scale_2`` that the + caller applies to the fp32 output.""" + try: + import gguf + except ImportError as err: + raise ImportError( + "Loading Flint NVFP4 GGUF checkpoints requires the `gguf` package " + "(pip install 'gguf>=0.19.0')." + ) from err + + reader = gguf.GGUFReader(path) + prec_a4 = reader.fields.get("general.tensor_extra.prec_a4") + if prec_a4 is not None and any(prec_a4.contents()): + raise ValueError( + "GGUF marks tensors with FP4 activations (prec_a4); only weight-only " + "NVFP4 (W4A16) is supported." + ) + tensors = {t.name: t for t in reader.tensors} + for name, t in tensors.items(): + if name.endswith(".scale"): + continue + data = _as_tensor(t.data) + if t.tensor_type != _GGUF_NVFP4: + hf_name, _ = _map_name(name) + if t.tensor_type == gguf.GGMLQuantizationType.BF16: + data = data.view(torch.bfloat16) + elif t.tensor_type != gguf.GGMLQuantizationType.F32: + raise ValueError( + f"Unsupported GGUF tensor type for {name}: {t.tensor_type}" + ) + yield hf_name, data + continue + + base = name.removesuffix(".weight") + hf_base, stacked = _map_name(base) + n_cols = int(t.shape[0]) + weight, weight_scale = _repack_nvfp4(data.to(device), n_cols) + global_scale = _as_tensor(tensors[base + ".scale"].data) + if hf_base in dequantize: + yield ( + hf_base + ".weight", + dequantize_nvfp4(weight, weight_scale, out_dtype=torch.bfloat16), + ) + yield hf_base + ".weight_scale_2", global_scale.reshape(()) + continue + if not stacked: + yield hf_base + ".weight", weight + yield hf_base + ".weight_scale", weight_scale + yield hf_base + ".weight_scale_2", global_scale.reshape(()) + continue + for e in range(weight.shape[0]): + expert = hf_base.format(e) + yield expert + ".weight", weight[e] + yield expert + ".weight_scale", weight_scale[e] + yield expert + ".weight_scale_2", global_scale[e] diff --git a/vllm/models/flint/model.py b/vllm/models/flint/model.py new file mode 100644 index 00000000000..8998fbfe851 --- /dev/null +++ b/vllm/models/flint/model.py @@ -0,0 +1,693 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-only Flint model (BF16 safetensors or NVFP4 GGUF checkpoints). + +Flint is a Qwen3-MoE-style decoder with learned attention sinks, hybrid +sliding-window/full attention (full layers use no RoPE), dense SwiGLU MLPs on +``mlp_only_layers``, sigmoid top-k routing with renormalized weights, and an +ungated shared expert. +""" + +from collections.abc import Iterable + +import torch +from torch import nn + +from vllm.config import ( + CacheConfig, + ParallelConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.attention.attention import ( + unified_attention_with_output, +) +from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import make_attention_sink +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors + +from .configs import FlintConfig +from .gguf_nvfp4 import find_gguf_file, gguf_nvfp4_weights_iterator +from .ops import linear as flint_linear +from .ops.moe import ( + BF16_DECODE_MAX_TOKENS, + LL_MOE_MAX_TOKENS, + FlintBf16DecodeMoE, + FlintNvfp4Experts, + LLMoEWeights, + flint_moe_ll, + ll_weights_from_layers, +) +from .ops.norm import RMSNorm, silu_and_mul +from .ops.qk_norm_rope import qk_norm_rope_cache +from .ops.qkv_fused import MAX_M as LL_MAX_TOKENS +from .ops.qkv_fused import flint_qkv_fused + + +def _flint_kernels( + quant_config: QuantizationConfig | None, parallel: ParallelConfig | None = None +) -> bool: + """Flint's kernels serve the BF16 and NVFP4 checkpoints on a single SM100 + GPU (TP=1, and for the MoE no EP/EPLB); other configurations use vLLM's + stock layers.""" + if not current_platform.is_device_capability_family(100): + return False + if get_tensor_model_parallel_world_size() != 1 or not ( + quant_config is None or isinstance(quant_config, ModelOptNvFp4Config) + ): + return False + return parallel is None or not ( + parallel.enable_expert_parallel or parallel.enable_eplb + ) + + +# Dequantized to BF16 (without the global scale, so exactly) at load so logits +# can be accumulated in fp32; the global scale is applied to the fp32 logits. +_GGUF_DEQUANTIZED = frozenset({"lm_head"}) + + +def _nvfp4_gguf_quant_config() -> QuantizationConfig: + # The GGUF reference runtime keeps activations in 16 bits (W4A16). + return ModelOptNvFp4Config( + quant_method="W4A16_NVFP4", + is_checkpoint_nvfp4_serialized=True, + exclude_modules=["model.embed_tokens", "*.mlp.gate", *_GGUF_DEQUANTIZED], + group_size=16, + ) + + +class FlintMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + quant_config: QuantizationConfig | None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + act = flint_linear.gate_up_silu_mul(self.gate_up_proj, x) + if act is None: + gate_up, _ = self.gate_up_proj(x) + act = silu_and_mul(gate_up) + x, _ = self.down_proj(act) + return x + + +class FlintSparseMoeBlock(nn.Module): + def __init__( + self, + config: FlintConfig, + quant_config: QuantizationConfig | None, + prefix: str = "", + ) -> None: + super().__init__() + self.top_k = config.num_experts_per_tok + self.gate = GateLinear( + config.hidden_size, + config.num_experts, + quant_config=None, + prefix=f"{prefix}.gate", + ) + self.shared_expert = FlintMLP( + hidden_size=config.hidden_size, + intermediate_size=config.shared_expert_intermediate_size, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_expert", + ) + vllm_config = get_current_vllm_config() + # NVFP4: Flint's MoE for all batch sizes. BF16: Flint's decode kernels + # for small batches on the TRTLLM-layout experts, TRTLLM MoE otherwise. + is_nvfp4 = isinstance(quant_config, ModelOptNvFp4Config) + use_flint = _flint_kernels(quant_config, vllm_config.parallel_config) + if use_flint: + assert config.router_score_function == "sigmoid" + assert config.norm_topk_prob + self.use_flint_kernel = use_flint and is_nvfp4 + self.bf16_decode: FlintBf16DecodeMoE | None = None + self._ll_weights: LLMoEWeights | None = None + if self.use_flint_kernel: + self.experts = FlintNvfp4Experts( + config.num_experts, + config.hidden_size, + config.moe_intermediate_size, + config.num_experts_per_tok, + vllm_config.scheduler_config.max_num_batched_tokens, + ) + return + self.experts = FusedMoEFactory( + shared_experts=self.shared_expert, + gate=self.gate, + num_experts=config.num_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + scoring_func=config.router_score_function, + quant_config=quant_config, + prefix=f"{prefix}.experts", + ) + if use_flint: + self.bf16_decode = FlintBf16DecodeMoE( + config.num_experts, + config.hidden_size, + config.moe_intermediate_size, + config.num_experts_per_tok, + ) + + def forward_ll( + self, attn_out: torch.Tensor, residual: torch.Tensor, norm: RMSNorm + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """``(mlp(norm(attn_out + residual)), attn_out + residual)`` via the + fused decode MoE block (NVFP4, M <= LL_MOE_MAX_TOKENS), else None.""" + if not self.use_flint_kernel or attn_out.shape[0] > LL_MOE_MAX_TOKENS: + return None + if self._ll_weights is None: + assert self.experts.weights is not None + self._ll_weights = ll_weights_from_layers( + self.gate.weight, + self.shared_expert.gate_up_proj, + self.shared_expert.down_proj, + self.experts.weights, + ) + return flint_moe_ll( + attn_out, + residual, + norm.weight, + norm.variance_epsilon, + self._ll_weights, + self.top_k, + self.experts.route_counter, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_flint_kernel: + return self.experts(hidden_states, self.gate.weight, self.shared_expert) + if ( + self.bf16_decode is not None + and hidden_states.shape[0] <= BF16_DECODE_MAX_TOKENS + ): + return self.bf16_decode( + hidden_states, + self.gate.weight, + self.experts.routed_experts, + self.shared_expert, + ) + return self.experts(hidden_states=hidden_states, router_logits=hidden_states) + + +class FlintAttention(nn.Module): + def __init__( + self, + config: FlintConfig, + layer_idx: int, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str = "", + ) -> None: + super().__init__() + tp_size = get_tensor_model_parallel_world_size() + self.head_dim = config.head_dim + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = config.num_key_value_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=config.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.rotary_emb = None + if layer_idx not in config.no_rope_layer_indices: + # Only its cos|sin table is used, by the fused kernels. Kept in + # fp32 like the reference implementation; get_rope memoizes, so + # all layers share one table. + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + rope_parameters=config.rope_parameters, + dtype=torch.float32, + ) + assert self.rotary_emb.is_neox_style + assert self.rotary_emb.rotary_dim == self.head_dim + is_sliding = config.layer_types[layer_idx] == "sliding_attention" + self.softmax_offset = make_attention_sink(self.num_heads, torch.float32) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.head_dim**-0.5, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=config.sliding_window if is_sliding else None, + prefix=f"{prefix}.attn", + sinks=self.softmax_offset, + ) + # The fused kernel writes the KV cache itself (in the FlashInfer layout) + # and then calls the attention backend directly. Imported here so the + # module imports on platforms without FlashInfer. + from vllm.v1.attention.backends.flashinfer import ( + FlashInferImpl, + FlashInferMetadata, + FlashInferTrtllmAPIDecode, + ) + + impl = self.attn.impl + assert isinstance(impl, FlashInferImpl), "Flint requires FlashInfer attention" + self._metadata_cls = FlashInferMetadata + self._decode_cls = FlashInferTrtllmAPIDecode + assert get_current_vllm_config().model_config.dtype == torch.bfloat16 + assert self.attn.kv_cache_dtype in ("auto", "bfloat16") + assert impl.kv_sharing_target_layer_name is None + self.impl = impl + # Flint decode attention and the fused low-latency decode path. + self.flint_decode = _flint_kernels(quant_config) + # Largest pure-decode batch served by the Flint kernel; beyond it + # FlashInfer is faster on the full-attention layers at long context. + self.max_decode_batch = 64 if is_sliding else 8 + if self.flint_decode: + assert not config.attention_bias + from .ops.decode_attention import flint_decode_attention + + self._decode_attention = flint_decode_attention + + def forward( + self, positions: torch.Tensor, hidden_states: torch.Tensor + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + k_cache, v_cache, layer_slots = self._cache_views() + plan = self._decode_plan() if self.flint_decode else None + q = qk_norm_rope_cache( + qkv, + positions, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.variance_epsilon, + self._cos_sin(), + k_cache, + v_cache, + layer_slots, + self.num_heads, + self.num_kv_heads, + self.head_dim, + # Flint decode attention waits on the grid dependency before + # reading q / K / V, so it may start its prologue early. + launch_dependents=plan is not None, + ) + num_tokens = qkv.shape[0] + attn_output = torch.empty_like(q) + if plan is not None: + self._run_decode_attention(q, attn_output, *plan) + else: + _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + unified_attention_with_output( + q.view(num_tokens, self.num_heads, self.head_dim), + k.view(num_tokens, self.num_kv_heads, self.head_dim), + v.view(num_tokens, self.num_kv_heads, self.head_dim), + attn_output.view(num_tokens, self.num_heads, self.head_dim), + self.attn.layer_name, + ) + output, _ = self.o_proj(attn_output) + return output + + def _cos_sin(self) -> torch.Tensor | None: + return None if self.rotary_emb is None else self.rotary_emb.cos_sin_cache + + def _cache_views( + self, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + """(k_cache, v_cache, slot_mapping) for the fused cache write, with the + caches as [blocks, block_size, kv_heads, head_dim] views; all None while + the KV cache is unallocated (memory profiling).""" + kv_cache = self.attn.kv_cache + if kv_cache.numel() == 0: + return None, None, None + # As in FlashInferImpl.do_kv_cache_update: (B, H, N, 2D) -> (B, N, H, D) + k_cache, v_cache = kv_cache.transpose(1, 2).split(self.head_dim, dim=-1) + slot_mapping = get_forward_context().slot_mapping + assert isinstance(slot_mapping, dict) + return k_cache, v_cache, slot_mapping.get(self.attn.layer_name) + + def _decode_plan(self) -> tuple | None: + """(decode metadata, num_decodes, window) when the batch is pure + single-token decode of at most ``max_decode_batch`` requests (the Flint + decode attention kernel), else None. K/V are already in the cache.""" + forward_metadata = get_forward_context().attn_metadata + attn_metadata = ( + forward_metadata.get(self.attn.layer_name) + if isinstance(forward_metadata, dict) + else forward_metadata + ) + if ( + not isinstance(attn_metadata, self._metadata_cls) # profiling run + or attn_metadata.num_prefills != 0 + or not 0 < attn_metadata.num_decodes <= self.max_decode_batch + or attn_metadata.num_decode_tokens != attn_metadata.num_decodes + ): + return None + decode = attn_metadata.decode + assert isinstance(decode, self._decode_cls) + assert decode.q_len_per_req == 1 and decode.dcp_query_start_loc is None + left = self.impl.window_left + window = left + 1 if left >= 0 else None + return decode, attn_metadata.num_decodes, window + + def _run_decode_attention( + self, q: torch.Tensor, out: torch.Tensor, decode, n: int, window: int | None + ) -> None: + shape = (-1, self.num_heads, self.head_dim) + self._decode_attention( + q.view(shape)[:n], + self.attn.kv_cache, + decode.block_tables, + decode.seq_lens, + self.impl.sinks, + window, + self.impl.scale, + out.view(shape)[:n], + ) + + def forward_ll( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + input_layernorm: RMSNorm, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Small pure-decode batches: two fused kernels for the residual-add + RMSNorm, qkv GEMV, q/k norm, RoPE and KV-cache write, then the Flint + decode attention. Returns None to fall back to the regular path.""" + if hidden_states.shape[0] > LL_MAX_TOKENS: + return None + k_cache, v_cache, layer_slots = self._cache_views() + plan = self._decode_plan() + if layer_slots is None or plan is None: + return None + layer = self.qkv_proj + bf16w = layer.weight.dtype == torch.bfloat16 + q, residual = flint_qkv_fused( + hidden_states, + residual, + input_layernorm.weight, + input_layernorm.variance_epsilon, + layer.weight, + None if bf16w else layer.flint_weight_scale, + None if bf16w else layer.weight_global_scale, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.variance_epsilon, + self._cos_sin(), + positions, + k_cache, + v_cache, + layer_slots, + self.num_heads, + self.num_kv_heads, + ) + attn_output = torch.empty_like(q) + self._run_decode_attention(q, attn_output, *plan) + output, _ = self.o_proj(attn_output) + return output, residual + + +class FlintDecoderLayer(nn.Module): + def __init__( + self, + config: FlintConfig, + cache_config: CacheConfig | None, + quant_config: QuantizationConfig | None, + prefix: str = "", + ) -> None: + super().__init__() + layer_idx = extract_layer_index(prefix) + self.self_attn = FlintAttention( + config, + layer_idx, + cache_config, + quant_config, + prefix=f"{prefix}.self_attn", + ) + self.mlp: FlintMLP | FlintSparseMoeBlock + if layer_idx in config.mlp_only_layers: + self.mlp = FlintMLP( + config.hidden_size, + config.intermediate_size, + quant_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = FlintSparseMoeBlock(config, quant_config, prefix=f"{prefix}.mlp") + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.self_attn.flint_decode: + ll = self.self_attn.forward_ll( + positions, hidden_states, residual, self.input_layernorm + ) + if ll is not None: + return self._post_attention(*ll) + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + hidden_states = self.self_attn(positions, hidden_states) + return self._post_attention(hidden_states, residual) + + def _post_attention( + self, hidden_states: torch.Tensor, residual: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if isinstance(self.mlp, FlintSparseMoeBlock): + ll = self.mlp.forward_ll( + hidden_states, residual, self.post_attention_layernorm + ) + if ll is not None: + return ll + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class FlintModel(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + quant_config: QuantizationConfig | None, + prefix: str = "", + ) -> None: + super().__init__() + config: FlintConfig = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + self.config = config + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: FlintDecoderLayer( + config, cache_config, quant_config, prefix=prefix + ), + prefix=f"{prefix}.layers", + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + else: + hidden_states = ( + inputs_embeds + if inputs_embeds is not None + else self.embed_input_ids(input_ids) + ) + residual = None + for layer in self.layers[self.start_layer : self.end_layer]: + hidden_states, residual = layer(positions, hidden_states, residual) + if self.end_layer != len(self.layers): + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class FlintForCausalLM(nn.Module, SupportsPP): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_stacked={ + ".q_proj": (".qkv_proj", "q"), + ".k_proj": (".qkv_proj", "k"), + ".v_proj": (".qkv_proj", "v"), + ".mlp.gate_proj": (".mlp.gate_up_proj", 0), + ".mlp.up_proj": (".mlp.gate_up_proj", 1), + ".shared_expert.gate_proj": (".shared_expert.gate_up_proj", 0), + ".shared_expert.up_proj": (".shared_expert.gate_up_proj", 1), + } + ) + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config: FlintConfig = vllm_config.model_config.hf_config + self.config = config + + model_config = vllm_config.model_config + self.gguf_path = find_gguf_file( + model_config.model, + model_config.revision, + vllm_config.load_config.download_dir, + ) + quant_config = vllm_config.quant_config + if self.gguf_path is not None: + if quant_config is not None: + raise ValueError( + "Flint NVFP4 GGUF checkpoints carry their own " + "quantization; do not pass --quantization." + ) + quant_config = _nvfp4_gguf_quant_config() + vllm_config.quant_config = quant_config + # load_weights reads the GGUF itself and never consumes the default + # loader's (lazy) iterator; the pattern just names the right file. + self.allow_patterns_overrides = ["*.gguf"] + if quant_config is not None: + quant_config.packed_modules_mapping = self.packed_modules_mapping + self.quant_config = quant_config + + self.model = FlintModel( + vllm_config, quant_config, prefix=maybe_prefix(prefix, "model") + ) + self.flint_dense = _flint_kernels(quant_config) + if self.flint_dense: + flint_linear.install_flint_dense_kernels(self.model) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + return self.logits_processor(self.lm_head, hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + if self.gguf_path is not None: + device = next(self.parameters()).device + weights = self._apply_dequantized_scales( + gguf_nvfp4_weights_iterator(self.gguf_path, device, _GGUF_DEQUANTIZED), + ) + if self.flint_dense: + weights = flint_linear.check_merged_global_scales(weights) + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def _apply_dequantized_scales( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + for name, tensor in weights: + if name == "lm_head.weight_scale_2": + self.logits_processor.scale = float(tensor) + else: + yield name, tensor diff --git a/vllm/models/flint/ops/__init__.py b/vllm/models/flint/ops/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm/models/flint/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/flint/ops/cute_primitives.py b/vllm/models/flint/ops/cute_primitives.py new file mode 100644 index 00000000000..2094a5e1f62 --- /dev/null +++ b/vllm/models/flint/ops/cute_primitives.py @@ -0,0 +1,428 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CuTe DSL device primitives for Flint's small-M BF16 x NVFP4/BF16 GEMV. + +Shared with other Flint kernels; keep signatures stable. A-fragment layout +(``mma.sync`` m16n8k16, swap-AB: weight rows are the MMA M dimension): lane +``(g, t) = (lane // 4, lane % 4)`` of k-tile ``j`` within a 64-wide K chunk +covers ``k = 64c + 16t + 4j + {0..3}`` of rows ``g`` and ``g + 8``. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op + + +def ptx(code, constraints, operands, results=None, *, effects=True, loc=None, ip=None): + if results is None: + ty = None + elif isinstance(results, tuple): + ty = llvm.StructType.get_literal([r.mlir_type for r in results]) + else: + ty = results.mlir_type + value = llvm.inline_asm( + ty, + [op.ir_value(loc=loc, ip=ip) for op in operands], + code, + constraints, + has_side_effects=effects, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + if results is None: + return None + if isinstance(results, tuple): + return tuple( + r(llvm.extractvalue(r.mlir_type, value, [i], loc=loc, ip=ip)) + for i, r in enumerate(results) + ) + return results(value) + + +@dsl_user_op +def smem_addr(ptr, *, loc=None, ip=None): + return Uint32(ptr.toint(loc=loc, ip=ip)) + + +# Loads are volatile so that every weight load is issued up front, in +# program order, before the dependent math. +@dsl_user_op +def ld_w8(addr: Int64, *, loc=None, ip=None): + return list( + ptx( + "ld.global.nc.L1::no_allocate.v2.u32 {$0, $1}, [$2];", + "=r,=r,l", + [addr], + (Uint32,) * 2, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_w8p(addr: Int64, *, loc=None, ip=None): + # ``ld_w8`` with a 128-byte L2 prefetch: a lane's later chunks of the same + # row follow in 32-byte steps. + return list( + ptx( + "ld.global.nc.L1::no_allocate.L2::128B.v2.u32 {$0, $1}, [$2];", + "=r,=r,l", + [addr], + (Uint32,) * 2, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_w16(addr: Int64, *, loc=None, ip=None): + return list( + ptx( + "ld.global.nc.L1::no_allocate.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + [addr], + (Uint32,) * 4, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def ld_u8(addr: Int64, *, loc=None, ip=None): + return ptx( + "{\n.reg .b16 h;\nld.global.nc.u8 h, [$1];\ncvt.u32.u16 $0, h;\n}", + "=r,l", + [addr], + Uint32, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_x16(addr: Int64, pred: Int32, *, loc=None, ip=None): + # 16 bf16 activations, or zeros for padded tokens. + return list( + ptx( + "{\n.reg .pred p;\nsetp.ne.b32 p, $9, 0;\n" + "mov.b32 $0, 0; mov.b32 $1, 0; mov.b32 $2, 0; mov.b32 $3, 0;\n" + "mov.b32 $4, 0; mov.b32 $5, 0; mov.b32 $6, 0; mov.b32 $7, 0;\n" + "@p ld.global.v4.u32 {$0, $1, $2, $3}, [$8];\n" + "@p ld.global.v4.u32 {$4, $5, $6, $7}, [$8+16];\n}", + "=r,=r,=r,=r,=r,=r,=r,=r,l,r", + [addr, pred], + (Uint32,) * 8, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def scale_pair(sc: Uint32, *, loc=None, ip=None): + # E4M3 byte -> BF16x2 (s, s); exact. + return ptx( + "{\n.reg .b16 h;\n.reg .b8 b;\ncvt.u16.u32 h, $1;\n" + "mov.b16 {b, _}, h;\nmov.b16 h, {b, b};\n" + "cvt.rn.bf16x2.e4m3x2 $0, h;\n}", + "=r,r", + [sc], + Uint32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def decode( + lo: Uint32, + hi: Uint32, + j: cutlass.Constexpr, + sg: Uint32, + sg8: Uint32, + *, + loc=None, + ip=None, +): + # A fragment of k-tile j: bytes (2j, 2j+1) % 4 of the row-g word `lo` + # and the row-(g+8) word `hi`, scaled exactly in BF16. + b = (2 * j) % 4 + return list( + ptx( + "{\n.reg .b8 b0, b1, b2, b3, c0, c1, c2, c3;\n" + "mov.b32 {b0, b1, b2, b3}, $4;\nmov.b32 {c0, c1, c2, c3}, $5;\n" + f"cvt.rn.bf16x2.e2m1x2 $0, b{b};\n" + f"cvt.rn.bf16x2.e2m1x2 $1, c{b};\n" + f"cvt.rn.bf16x2.e2m1x2 $2, b{b + 1};\n" + f"cvt.rn.bf16x2.e2m1x2 $3, c{b + 1};\n" + "mul.rn.bf16x2 $0, $0, $6;\nmul.rn.bf16x2 $1, $1, $7;\n" + "mul.rn.bf16x2 $2, $2, $6;\nmul.rn.bf16x2 $3, $3, $7;\n}", + "=r,=r,=r,=r,r,r,r,r", + [lo, hi, sg, sg8], + (Uint32,) * 4, + effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def mma(c, a, b0: Uint32, b1: Uint32, *, loc=None, ip=None): + return list( + ptx( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{$0, $1, $2, $3}, {$8, $9, $10, $11}, {$12, $13}, {$4, $5, $6, $7};", + "=f,=f,=f,=f,f,f,f,f,r,r,r,r,r,r", + [*c, *a, b0, b1], + (Float32,) * 4, + effects=False, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def st_out(addr: Int64, v: Float32, pred: Int32, *, loc=None, ip=None): + ptx( + "{\n.reg .pred p;\n.reg .b16 h;\nsetp.ne.b32 p, $2, 0;\n" + "cvt.rn.bf16.f32 h, $1;\n@p st.global.b16 [$0], h;\n}", + "l,f,r", + [addr, v, pred], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def st_dsmem_f32(smem_ptr, v: Float32, *, loc=None, ip=None): + remote = cute.arch.map_dsmem_ptr(smem_ptr, Int32(0), loc=loc, ip=ip) + addr = Int32(remote.toint(loc=loc, ip=ip)) + ptx("st.shared::cluster.f32 [$0], $1;", "r,f", [addr, v], loc=loc, ip=ip) + + +@dsl_user_op +def silu_mul(g: Float32, u: Float32, *, loc=None, ip=None): + # g / (1 + expf(-g)) * u in fp32 (rounded once by the caller's store), with + # CUDA's precise expf (libdevice __nv_expf) and an IEEE division. + return ptx( + "{\n.reg .b32 x, t, j, f, e, sc, d, a;\n" + "neg.f32 x, $1;\n" + "fma.rn.ftz.f32 t, x, 0f3BBB989D, 0f3F000000;\n" + "cvt.ftz.sat.f32.f32 t, t;\n" + "fma.rm.ftz.f32 j, t, 0f437C0000, 0f4B400001;\n" + "add.f32 f, j, 0fCB40007F;\n" + "neg.f32 f, f;\n" + "fma.rn.ftz.f32 f, x, 0f3FB8AA3B, f;\n" + "fma.rn.ftz.f32 f, x, 0f32A57060, f;\n" + "shl.b32 sc, j, 23;\n" + "ex2.approx.ftz.f32 e, f;\n" + "fma.rn.f32 d, e, sc, 0f3F800000;\n" + "div.rn.f32 a, $1, d;\n" + "mul.f32 $0, a, $2;\n}", + "=f,f,f", + [g, u], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def round_bf16(v: Float32, *, loc=None, ip=None): + return ptx( + "{\n.reg .b16 h;\ncvt.rn.bf16.f32 h, $1;\ncvt.f32.bf16 $0, h;\n}", + "=f,f", + [v], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def sts_f32(addr, a, *, loc=None, ip=None): + ptx( + "st.shared.f32 [$0], $1;", + "r,f,~{memory}", + [Uint32(addr), Float32(a)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def lds_f32(addr, *, loc=None, ip=None): + return ptx( + "ld.shared.f32 $0, [$1];", + "=f,r,~{memory}", + [Uint32(addr)], + Float32, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg_bf16(addr, a, *, loc=None, ip=None): + ptx( + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; st.global.b16 [$0], h; }", + "l,f,~{memory}", + [Int64(addr), Float32(a)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def stg_u32x4(addr, a, b, c, d, *, loc=None, ip=None): + ptx( + "st.global.v4.b32 [$0], {$1,$2,$3,$4};", + "l,r,r,r,r,~{memory}", + [Int64(addr), Uint32(a), Uint32(b), Uint32(c), Uint32(d)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def bf16x2(lo, hi, *, loc=None, ip=None): + return ptx( + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + [Float32(hi), Float32(lo)], + Uint32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def fma(a, b, c, *, loc=None, ip=None): + return ptx( + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + [Float32(a), Float32(b), Float32(c)], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + +@cute.jit +def load_rows( + w_base, + s_base, + row, + KC: cutlass.Constexpr, + CPW: cutlass.Constexpr, + c_begin, + t, + qs, + ss, + q_off: cutlass.Constexpr, + s_off: cutlass.Constexpr, + PREFETCH: cutlass.Constexpr = False, +): + """NVFP4 weight words / scale bytes of rows (row, row + 8) for CPW chunks + from ``c_begin`` (ModelOpt weight, 128x4-swizzled scales).""" + kb = KC * 32 + wbase = w_base + Int64(row) * kb + c_begin * 32 + t * 8 + row8 = row + 8 + sb_g = ( + s_base + + Int64(row // 128) * (KC * 512) + + (row % 32) * 16 + + ((row % 128) // 32) * 4 + + t + + c_begin * 512 + ) + sb_g8 = ( + s_base + + Int64(row8 // 128) * (KC * 512) + + (row8 % 32) * 16 + + ((row8 % 128) // 32) * 4 + + t + + c_begin * 512 + ) + for c in cutlass.range_constexpr(CPW): + if cutlass.const_expr(PREFETCH): + qa = ld_w8p(wbase + c * 32) + qb = ld_w8p(wbase + 8 * kb + c * 32) + else: + qa = ld_w8(wbase + c * 32) + qb = ld_w8(wbase + 8 * kb + c * 32) + q0 = q_off + c * 4 + qs[q0 + 0] = qa[0] + qs[q0 + 1] = qa[1] + qs[q0 + 2] = qb[0] + qs[q0 + 3] = qb[1] + ss[s_off + c * 2] = ld_u8(sb_g + c * 512) + ss[s_off + c * 2 + 1] = ld_u8(sb_g8 + c * 512) + + +def alloc_reduce(smem, NA, WARPS, SPLIT): + red = smem.allocate_tensor( + Float32, + cute.make_layout((WARPS, NA, 32), stride=(NA * 32, 32, 1)), + byte_alignment=16, + ) + part = smem.allocate_tensor( + Float32, + cute.make_layout((SPLIT, NA, 32), stride=(NA * 32, 32, 1)), + byte_alignment=16, + ) + return red, part + + +@cute.jit +def cta_reduce( + acc, + NA: cutlass.Constexpr, + WARPS: cutlass.Constexpr, + SPLIT: cutlass.Constexpr, + warp, + lane, + rank, + red, + part, +): + """Sum each lane's NA accumulators over the CTA's warps and the cluster's + CTAs (fixed order) into warp 0 of cluster rank 0. With SPLIT > 1 the + caller must have issued ``cluster_arrive_relaxed``.""" + for i in cutlass.range_constexpr(NA): + red[warp, i, lane] = acc[i] + cute.arch.sync_threads() + if cutlass.const_expr(SPLIT > 1): + cute.arch.cluster_wait() + if warp == 0: + for i in cutlass.range_constexpr(NA): + v = acc[i] + for w2 in cutlass.range_constexpr(1, WARPS): + v = v + red[w2, i, lane] + acc[i] = v + if cutlass.const_expr(SPLIT > 1): + for i in cutlass.range_constexpr(NA): + st_dsmem_f32(part.iterator + (rank * NA + i) * 32 + lane, acc[i]) + if cutlass.const_expr(SPLIT > 1): + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + if warp == 0 and rank == 0: + for i in cutlass.range_constexpr(NA): + v = part[0, i, lane] + for r2 in cutlass.range_constexpr(1, SPLIT): + v = v + part[r2, i, lane] + acc[i] = v diff --git a/vllm/models/flint/ops/decode_attention.py b/vllm/models/flint/ops/decode_attention.py new file mode 100644 index 00000000000..d00409fd828 --- /dev/null +++ b/vllm/models/flint/ops/decode_attention.py @@ -0,0 +1,1134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Split-KV paged decode attention for Flint (GQA 8, head_dim 128, bf16 KV). + +Numerics mirror FlashInfer's trtllm-gen decode kernel: ``S = K q`` with bf16 +inputs and fp32 accumulation; softmax in fp32 as ``exp2(S * c - m * c)`` with +``c = scale * log2(e)`` and the running max ``m`` in unscaled logits; the row +sum accumulates the fp32 probabilities while ``P`` is rounded to bf16 for +``P V`` (fp32 accumulation); the sink adds ``exp2(sink * log2(e) - m * c)`` to +the final denominator. Split partials (unnormalized O, m, l) stay fp32 +(trtllm-gen stores partial O in bf16) and are merged either by the last CTA of +each (request, kv head), found with a self-resetting arrival counter, or, for +the short sliding-window layers, inside a thread-block cluster through +distributed shared memory (every rank merges a slice of the head dims). + +Each warp streams 8 KB pages (16 tokens x [K | V], HND layout) through a +private ring of bulk-copy (TMA) stages and runs ``mma.sync`` on them. The +kernel is launched with programmatic dependent launch: it prefetches its +leading pages before waiting on the grid dependency (all but the page holding +the current token, whose K/V the previous kernel writes), touches q and its +workspace only after the wait, and triggers its dependents once K/V streaming +is done. +""" + +from functools import cache + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass.cutlass_dsl import dsl_user_op +from quack.compile_utils import make_fake_tensor + +from .cute_primitives import ( + bf16x2, + fma, + lds_f32, + mma, + ptx, + smem_addr, + stg_bf16, + stg_u32x4, + sts_f32, +) + +_PAGE = 16 +_GQA = 8 +_HEAD_DIM = 128 +_PAGE_BYTES = _PAGE * 2 * _HEAD_DIM * 2 +_ROW_BYTES = 2 * _HEAD_DIM * 2 +# Per-warp CTA-combine buffer (fp32): O[8][128], m[8], l[8]. +_WBUF = _GQA * _HEAD_DIM + 2 * _GQA +_LOG2E = 1.4426950408889634 +_SMEM_PER_SM = 228 * 1024 + + +@dsl_user_op +def _bulk_copy(dst, src, mbar, nbytes: int, *, loc=None, ip=None): + ptx( + "{ .reg .b64 p; createpolicy.fractional.L2::evict_first.b64 p, 1.0; " + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + f".L2::cache_hint [$0], [$1], {nbytes}, [$2], p; }}", + "r,l,r,~{memory}", + [Uint32(dst), Int64(src), Uint32(mbar)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds128(addr, *, loc=None, ip=None): + return ptx( + "ld.shared.v4.b32 {$0,$1,$2,$3}, [$4];", + "=r,=r,=r,=r,r,~{memory}", + [Uint32(addr)], + (Uint32,) * 4, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldg128(addr, *, loc=None, ip=None): + return ptx( + "ld.global.nc.v4.b32 {$0,$1,$2,$3}, [$4];", + "=r,=r,=r,=r,l,~{memory}", + [Int64(addr)], + (Uint32,) * 4, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldg_cg_f32x4(addr, *, loc=None, ip=None): + return ptx( + "ld.global.cg.v4.f32 {$0,$1,$2,$3}, [$4];", + "=f,=f,=f,=f,l,~{memory}", + [Int64(addr)], + (Float32,) * 4, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ldg_i32_pred(addr, pred, *, loc=None, ip=None): + return ptx( + "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b32 $0, 0; " + "@p ld.global.nc.b32 $0, [$1]; }", + "=r,l,r", + [Int64(addr), Int32(pred)], + Int32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_f32x4(addr, a, b, c, d, *, loc=None, ip=None): + ptx( + "st.global.v4.f32 [$0], {$1,$2,$3,$4};", + "l,f,f,f,f,~{memory}", + [Int64(addr), Float32(a), Float32(b), Float32(c), Float32(d)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_f32(addr, a, *, loc=None, ip=None): + ptx( + "st.global.f32 [$0], $1;", + "l,f,~{memory}", + [Int64(addr), Float32(a)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u32x2(addr, a, b, *, loc=None, ip=None): + ptx( + "st.global.v2.b32 [$0], {$1,$2};", + "l,r,r,~{memory}", + [Int64(addr), Uint32(a), Uint32(b)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _sts_i32(addr, a, *, loc=None, ip=None): + ptx( + "st.shared.b32 [$0], $1;", + "r,r,~{memory}", + [Uint32(addr), Int32(a)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _lds_i32(addr, *, loc=None, ip=None): + return ptx( + "ld.shared.b32 $0, [$1];", + "=r,r,~{memory}", + [Uint32(addr)], + Int32, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _atomic_add_acq_rel(addr, val, *, loc=None, ip=None): + return ptx( + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", + "=r,l,r,~{memory}", + [Int64(addr), Int32(val)], + Int32, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_i32_relaxed(addr, val, *, loc=None, ip=None): + ptx( + "st.relaxed.gpu.global.b32 [$0], $1;", + "l,r,~{memory}", + [Int64(addr), Int32(val)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _mapa(addr, rank, *, loc=None, ip=None): + return ptx( + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + [Uint32(addr), Int32(rank)], + Uint32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_cluster_f32x4(addr, a, b, c, d, *, loc=None, ip=None): + ptx( + "st.shared::cluster.v4.f32 [$0], {$1,$2,$3,$4};", + "r,f,f,f,f,~{memory}", + [Uint32(addr), Float32(a), Float32(b), Float32(c), Float32(d)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _stg_u32(addr, a, *, loc=None, ip=None): + ptx( + "st.global.b32 [$0], $1;", + "l,r,~{memory}", + [Int64(addr), Uint32(a)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _st_cluster_f32(addr, a, *, loc=None, ip=None): + ptx( + "st.shared::cluster.f32 [$0], $1;", + "r,f,~{memory}", + [Uint32(addr), Float32(a)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _prmt(a, b, sel, *, loc=None, ip=None): + return ptx( + "prmt.b32 $0, $1, $2, $3;", + "=r,r,r,r", + [Uint32(a), Uint32(b), Uint32(sel)], + Uint32, + effects=False, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _ex2(x, *, loc=None, ip=None): + return ptx( + "ex2.approx.ftz.f32 $0, $1;", + "=f,f", + [Float32(x)], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + +def _neg_inf(): + return Float32(float("-inf")) + + +class _FlintDecodeAttention: + """Grid (splits, kv heads, batch). The attended pages of a (request, kv + head) are interleaved over the warps of its active splits; a split is + active when every one of its warps gets at least ``min_pages`` pages. + With ``cluster`` the splits of a (request, kv head) form a thread-block + cluster and merge through distributed shared memory instead of global + partials.""" + + def __init__( + self, + num_kv_heads, + num_splits, + num_warps, + num_stages, + sliding, + min_pages, + cluster, + ): + self.num_kv_heads = num_kv_heads + self.num_splits = num_splits + self.num_warps = num_warps + self.num_stages = num_stages + self.sliding = sliding + self.min_pages = min_pages + self.cluster = cluster + self.ctas_per_sm = 0 if sliding else max(1, 2 * 4 // num_warps) + + @cute.jit + def __call__( + self, + q, + kv, + bt, + seq_lens, + sinks, + out, + part_o, + part_ml, + counters, + scale_log2: Float32, + window: Int32, + stream, + ): + batch = seq_lens.shape[0] + self.kernel( + q, + kv, + bt, + seq_lens, + sinks, + out, + part_o, + part_ml, + counters, + scale_log2, + window, + ).launch( + grid=(self.num_splits, self.num_kv_heads, batch), + block=(32 * self.num_warps, 1, 1), + stream=stream, + cluster=(self.num_splits, 1, 1) if self.cluster else None, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + q, + kv, + bt, + seq_lens, + sinks, + out, + part_o, + part_ml, + counters, + scale_log2: Float32, + window: Int32, + ): + NW = self.num_warps + NS = self.num_stages + NSPLIT = self.num_splits + H = self.num_kv_heads + MINP = self.min_pages + tid, _, _ = cute.arch.thread_idx() + split, kvh, b = cute.arch.block_idx() + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane = cute.arch.lane_idx() + g = lane >> 2 + t = lane & 3 + + smem = cutlass.memory.SmemAllocator() + ring = smem.allocate_array( + cutlass.Uint8, NW * NS * _PAGE_BYTES, byte_alignment=128 + ) + if cutlass.const_expr(self.ctas_per_sm > 0): + # Pad shared memory so at most ``ctas_per_sm`` CTAs share an SM: + # with PDL the next grid's CTAs launch while this one drains and + # could otherwise pile up on the first SMs to free up. + pad = _SMEM_PER_SM // (self.ctas_per_sm + 1) + 1 - NW * NS * _PAGE_BYTES + if cutlass.const_expr(pad > 0): + smem.allocate_array(cutlass.Uint8, pad, byte_alignment=128) + ring_base = smem_addr(ring) + mbars = smem.allocate_array(Int64, NW * NS, byte_alignment=8) + flag = smem.allocate_array(Int32, 4, byte_alignment=16) + flag_addr = smem_addr(flag) + cbuf_addr = flag_addr + if cutlass.const_expr(self.cluster): + # Per rank: this rank's 128 / NSPLIT dims of every split's partial + # O ([split][head][dim]), then every split's m and l per head. + cbuf = smem.allocate_array( + Float32, + NSPLIT * 8 * (_HEAD_DIM // NSPLIT) + 16 * NSPLIT, + byte_alignment=16, + ) + cbuf_addr = smem_addr(cbuf) + + # seq_lens / block_tables are read before griddepcontrol.wait by design: + # they are written by host-to-device copies ahead of the forward, not by + # the preceding kernel (only q, the current K/V and the workspace are). + seq_len = seq_lens[b] + start_tok = Int32(0) + if cutlass.const_expr(self.sliding): + start_tok = cutlass.max(seq_len - window, Int32(0)) + start_page = start_tok // _PAGE + num_pages = (seq_len + _PAGE - 1) // _PAGE - start_page + num_splits = cutlass.min( + Int32(NSPLIT), (num_pages + NW * MINP - 1) // (NW * MINP) + ) + if cutlass.const_expr(self.cluster): + # Every rank of the cluster merges a slice of the output. + num_splits = Int32(NSPLIT) + + bh = b * H + kvh + q_base = q.iterator.toint() + out_base = out.iterator.toint() + kv_base = kv.iterator.toint() + cnt_base = counters.iterator.toint() + arrive_addr = cnt_base + Int64(bh) * 4 + sq0 = Int64(q.stride[0]) + sq1 = Int64(q.stride[1]) + so0 = Int64(out.stride[0]) + so1 = Int64(out.stride[1]) + skb = Int64(kv.stride[0]) + skh = Int64(kv.stride[1]) + sbt = Int64(bt.stride[0]) + + # CTA epilogue mapping: E consecutive dims of one head per thread. + E = 1024 // (32 * NW) + TPH = 128 // E + eh = tid // TPH + ed = (tid % TPH) * E + out_row = out_base + (Int64(b) * so0 + Int64(kvh * _GQA + eh) * so1) * 2 + + if num_pages <= 0: + cute.arch.griddepcontrol_wait() + if split == 0: + if cutlass.const_expr(E == 4): + _stg_u32x2(out_row + ed * 2, 0, 0) + for i in cutlass.range_constexpr(E // 8): + stg_u32x4(out_row + (ed + 8 * i) * 2, 0, 0, 0, 0) + elif split < num_splits: + # Pages are interleaved over all warps of this (request, kv head). + WH = num_splits * NW + gw = split * NW + warp + n_w = (cutlass.max(num_pages - gw, Int32(0)) + WH - 1) // WH + bt_head = ( + bt.iterator.toint() + (Int64(b) * sbt + Int64(start_page + gw)) * 4 + ) + kv_head = kv_base + Int64(kvh) * skh * 2 + win_cur = _ldg_i32_pred(bt_head + Int64(WH * lane) * 4, Int32(lane < n_w)) + win_nxt = _ldg_i32_pred( + bt_head + Int64(WH * (lane + 32)) * 4, Int32(lane + 32 < n_w) + ) + warp_ring = ring_base + Uint32(warp * NS * _PAGE_BYTES) + warp_mbar = mbars + warp * NS + if lane == 0: + for s in cutlass.range_constexpr(NS): + cute.arch.mbarrier_init(warp_mbar + s, 1) + cute.arch.mbarrier_init_fence() + cute.arch.sync_warp() + if cutlass.const_expr(self.cluster): + # Peers must be running before DSMEM stores (waited on there). + cute.arch.cluster_arrive_relaxed() + + # Only the page holding the current token (its K/V write), q and + # the workspace depend on the previous kernel: prefetch the other + # leading pages before waiting on it. + p_last = (seq_len - 1) // _PAGE + for s in cutlass.range_constexpr(NS - 1): + if (s < n_w) & (start_page + gw + WH * s != p_last): + self.issue( + kv_head, + skb, + cute.arch.shuffle_sync(win_cur, s), + warp_ring + Uint32(s * _PAGE_BYTES), + warp_mbar + s, + lane, + ) + cute.arch.griddepcontrol_wait() + + # Q as the B operand of S^T = K Q^T: lane (g, t) holds head g, dims + # 32i + 8t .. +7. + qf = cute.make_rmem_tensor((16,), Uint32) + q_row = q_base + (Int64(b) * sq0 + Int64(kvh * _GQA + g) * sq1) * 2 + for i in cutlass.range_constexpr(4): + w = _ldg128(q_row + (32 * i + 8 * t) * 2) + for j in cutlass.range_constexpr(4): + qf[4 * i + j] = w[j] + + m0 = _neg_inf() + m1 = _neg_inf() + l0 = Float32(0.0) + l1 = Float32(0.0) + acc = cute.make_rmem_tensor((32,), Float32) + for i in cutlass.range_constexpr(32): + acc[i] = Float32(0.0) + kv_regs = cute.make_rmem_tensor((64,), Uint32) + + for s in cutlass.range_constexpr(NS - 1): + if (s < n_w) & (start_page + gw + WH * s == p_last): + self.issue( + kv_head, + skb, + cute.arch.shuffle_sync(win_cur, s), + warp_ring + Uint32(s * _PAGE_BYTES), + warp_mbar + s, + lane, + ) + for k in cutlass.range(n_w, unroll=1): + # Refill the slot consumed in the previous iteration. + kn = k + NS - 1 + if (kn & 31) == 0: + win_cur = win_nxt + win_nxt = _ldg_i32_pred( + bt_head + Int64(WH * (kn + 32 + lane)) * 4, + Int32(kn + 32 + lane < n_w), + ) + # Every lane's shared-memory reads of this slot (previous + # iteration) are done before lane 0 refills it. + cute.arch.sync_warp() + if kn < n_w: + self.issue( + kv_head, + skb, + cute.arch.shuffle_sync(win_cur, kn & 31), + warp_ring + Uint32((kn % NS) * _PAGE_BYTES), + warp_mbar + kn % NS, + lane, + ) + slot = k % NS + cute.arch.mbarrier_wait(warp_mbar + slot, (k // NS) & 1) + page = start_page + gw + WH * k + self.load_page( + kv_regs, + warp_ring + Uint32(slot * _PAGE_BYTES), + page, + start_tok, + seq_len, + g, + t, + ) + m0, m1, l0, l1 = self.compute( + kv_regs, + qf, + acc, + page, + start_tok, + seq_len, + scale_log2, + m0, + m1, + l0, + l1, + g, + t, + lane, + ) + + for off in cutlass.range_constexpr(3): + l0 = l0 + cute.arch.shuffle_sync_bfly(l0, 4 << off) + l1 = l1 + cute.arch.shuffle_sync_bfly(l1, 4 << off) + # K/V streaming is done: let a PDL-launched successor start its + # prologue (e.g. o_proj weight loads) during the combine. + cute.arch.griddepcontrol_launch_dependents() + + # CTA combine through shared memory (reuses the ring). + cute.arch.sync_threads() + wbuf = ring_base + Uint32(warp * _WBUF * 4) + for j in cutlass.range_constexpr(8): + d0 = 8 * g + 2 * j if j < 4 else 64 + 8 * g + 2 * (j - 4) + sts_f32(wbuf + ((2 * t) * 128 + d0) * 4, acc[4 * j]) + sts_f32(wbuf + ((2 * t + 1) * 128 + d0) * 4, acc[4 * j + 1]) + sts_f32(wbuf + ((2 * t) * 128 + d0 + 1) * 4, acc[4 * j + 2]) + sts_f32(wbuf + ((2 * t + 1) * 128 + d0 + 1) * 4, acc[4 * j + 3]) + if g == 0: + sts_f32(wbuf + (1024 + 2 * t) * 4, m0) + sts_f32(wbuf + (1024 + 2 * t + 1) * 4, m1) + sts_f32(wbuf + (1032 + 2 * t) * 4, l0) + sts_f32(wbuf + (1032 + 2 * t + 1) * 4, l1) + cute.arch.sync_threads() + + m_cta = _neg_inf() + for w in cutlass.range_constexpr(NW): + m_cta = cute.arch.fmax( + m_cta, lds_f32(ring_base + (w * _WBUF + 1024 + eh) * 4) + ) + m_cta_s = m_cta * scale_log2 + l_cta = Float32(0.0) + o = cute.make_rmem_tensor((E,), Float32) + for e in cutlass.range_constexpr(E): + o[e] = Float32(0.0) + for w in cutlass.range_constexpr(NW): + mw = lds_f32(ring_base + (w * _WBUF + 1024 + eh) * 4) + corr = Float32(0.0) + if mw != _neg_inf(): + corr = _ex2(fma(mw, scale_log2, -m_cta_s)) + l_cta = l_cta + corr * lds_f32(ring_base + (w * _WBUF + 1032 + eh) * 4) + for e in cutlass.range_constexpr(E): + o[e] = o[e] + corr * lds_f32( + ring_base + (w * _WBUF + eh * 128 + ed + e) * 4 + ) + + sink_s = Float32(sinks[kvh * _GQA + eh]) * _LOG2E + if cutlass.const_expr(self.cluster): + self.cluster_reduce( + cbuf_addr, split, tid, eh, ed, o, m_cta, l_cta, sinks, kvh, + scale_log2, out_base + Int64(b) * so0 * 2, so1, E, + ) # fmt: skip + elif num_splits == 1: + inv = Float32(1.0) / (l_cta + _ex2(sink_s - m_cta_s)) + self.store_out(out_row, ed, o, inv, E) + else: + prow = Int64(bh * NSPLIT + split) + po = part_o.iterator.toint() + (prow * 1024 + eh * 128 + ed) * 4 + for i in cutlass.range_constexpr(E // 4): + _stg_f32x4( + po + 16 * i, o[4 * i], o[4 * i + 1], o[4 * i + 2], o[4 * i + 3] + ) + pml = part_ml.iterator.toint() + prow * 64 + if ed == 0: + _stg_f32(pml + eh * 4, m_cta) + _stg_f32(pml + (8 + eh) * 4, l_cta) + cute.arch.sync_threads() + if tid == 0: + _sts_i32(flag_addr, _atomic_add_acq_rel(arrive_addr, 1)) + cute.arch.sync_threads() + if _lds_i32(flag_addr) == num_splits - 1: + self.combine_splits( + part_o, + part_ml, + bh, + num_splits, + eh, + ed, + tid, + ring_base + Uint32(NW * _WBUF * 4), + scale_log2, + sink_s, + out_row, + o, + E, + ) + if tid == 0: + _stg_i32_relaxed(arrive_addr, 0) + + @cute.jit + def cluster_reduce( + self, cbuf_addr, split, tid, eh, ed, o, m_cta, l_cta, sinks, kvh, + scale_log2, out_b, so1, E: cutlass.Constexpr, + ): # fmt: skip + """Distributed cluster merge: each CTA sends rank r the dims + [r * DPR, (r + 1) * DPR) of its partial O (plus m, l); after one + cluster barrier every rank merges all splits for its dims.""" + NSPLIT = self.num_splits + NW = self.num_warps + DPR = _HEAD_DIM // NSPLIT + M_OFF = NSPLIT * 8 * DPR + L_OFF = M_OFF + 8 * NSPLIT + rank = ed // DPR + peer = _mapa(cbuf_addr, rank) + cute.arch.cluster_wait() + dst = peer + ((split * 8 + eh) * DPR + ed % DPR) * 4 + for i in cutlass.range_constexpr(E // 4): + _st_cluster_f32x4( + dst + 16 * i, o[4 * i], o[4 * i + 1], o[4 * i + 2], o[4 * i + 3] + ) + if ed % DPR == 0: + _st_cluster_f32(peer + (M_OFF + split * 8 + eh) * 4, m_cta) + _st_cluster_f32(peer + (L_OFF + split * 8 + eh) * 4, l_cta) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # This rank owns dims [split * DPR, (split + 1) * DPR) of all heads: + # OPT consecutive dims of one head per thread. + OPT = 8 * DPR // (32 * NW) + h = tid * OPT // DPR + dl = tid * OPT % DPR + m_all = _neg_inf() + for r in cutlass.range_constexpr(NSPLIT): + m_all = cute.arch.fmax(m_all, lds_f32(cbuf_addr + (M_OFF + r * 8 + h) * 4)) + m_all_s = m_all * scale_log2 + l_all = Float32(0.0) + oo = [Float32(0.0)] * OPT + for r in cutlass.range_constexpr(NSPLIT): + mr = lds_f32(cbuf_addr + (M_OFF + r * 8 + h) * 4) + corr = Float32(0.0) + if mr != _neg_inf(): + corr = _ex2(fma(mr, scale_log2, -m_all_s)) + l_all = l_all + corr * lds_f32(cbuf_addr + (L_OFF + r * 8 + h) * 4) + for e in cutlass.range_constexpr(OPT): + oo[e] = oo[e] + corr * lds_f32( + cbuf_addr + ((r * 8 + h) * DPR + dl + e) * 4 + ) + sink_s = Float32(sinks[kvh * _GQA + h]) * _LOG2E + inv = Float32(1.0) / (l_all + _ex2(sink_s - m_all_s)) + dst_g = out_b + (Int64(kvh * _GQA + h) * so1 + split * DPR + dl) * 2 + for e in cutlass.range_constexpr(0, OPT, 2): + if cutlass.const_expr(e + 1 < OPT): + _stg_u32(dst_g + e * 2, bf16x2(oo[e] * inv, oo[e + 1] * inv)) + else: + stg_bf16(dst_g + e * 2, oo[e] * inv) + + @cute.jit + def combine_splits( + self, + part_o, + part_ml, + bh, + num_splits, + eh, + ed, + tid, + ml_smem, + scale_log2, + sink_s, + out_row, + o, + E: cutlass.Constexpr, + ): + """Last CTA of a (request, kv head): merge the fp32 split partials.""" + NSPLIT = self.num_splits + ml_base = part_ml.iterator.toint() + Int64(bh * NSPLIT) * 64 + o_base = ( + part_o.iterator.toint() + (Int64(bh * NSPLIT) * 1024 + eh * 128 + ed) * 4 + ) + # Rounds of RCH splits keep at most ~40 vector loads per thread in flight. + RCH = min(NSPLIT, max(1, 160 // E)) + vs = [None] * RCH + for r in cutlass.range_constexpr(RCH): + sr = Int64(cutlass.min(Int32(r), num_splits - 1)) + vs[r] = [_ldg_cg_f32x4(o_base + sr * 4096 + 16 * i) for i in range(E // 4)] + for i in cutlass.range_constexpr( + (NSPLIT * 4 + 32 * self.num_warps - 1) // (32 * self.num_warps) + ): + idx = tid + i * 32 * self.num_warps + if idx < num_splits * 4: + v = _ldg_cg_f32x4(ml_base + Int64(idx) * 16) + for j in cutlass.range_constexpr(4): + sts_f32(ml_smem + (idx * 4 + j) * 4, v[j]) + cute.arch.sync_threads() + # Warp w owns heads w, w + NW, ..; its lanes cover the splits. The split + # weights replace l in ml_smem, the per-head scaled max and sum follow. + warp = tid // 32 + lane = tid % 32 + for hh in cutlass.range_constexpr(_GQA // self.num_warps): + h = warp + hh * self.num_warps + NL = (NSPLIT + 31) // 32 + mloc = [_neg_inf()] * NL + for r in cutlass.range_constexpr(NL): + sidx = lane + 32 * r + if sidx < num_splits: + mloc[r] = lds_f32(ml_smem + (sidx * 16 + h) * 4) + mh = mloc[0] + for r in cutlass.range_constexpr(1, NL): + mh = cute.arch.fmax(mh, mloc[r]) + for off in cutlass.range_constexpr(5): + mh = cute.arch.fmax(mh, cute.arch.shuffle_sync_bfly(mh, 1 << off)) + mh_s = mh * scale_log2 + lh = Float32(0.0) + for r in cutlass.range_constexpr(NL): + sidx = lane + 32 * r + if sidx < num_splits: + corr = _ex2(fma(mloc[r], scale_log2, -mh_s)) + lh = lh + corr * lds_f32(ml_smem + (sidx * 16 + 8 + h) * 4) + sts_f32(ml_smem + (sidx * 16 + 8 + h) * 4, corr) + for off in cutlass.range_constexpr(5): + lh = lh + cute.arch.shuffle_sync_bfly(lh, 1 << off) + if lane == 0: + sts_f32(ml_smem + (NSPLIT * 16 + h) * 4, mh_s) + sts_f32(ml_smem + (NSPLIT * 16 + 8 + h) * 4, lh) + cute.arch.sync_threads() + m_all_s = lds_f32(ml_smem + (NSPLIT * 16 + eh) * 4) + l_all = lds_f32(ml_smem + (NSPLIT * 16 + 8 + eh) * 4) + for e in cutlass.range_constexpr(E): + o[e] = Float32(0.0) + for r0 in cutlass.range_constexpr(0, NSPLIT, RCH): + if cutlass.const_expr(r0 > 0): # noqa: SIM102 (static vs dynamic) + if r0 < num_splits: + for r in cutlass.range_constexpr(RCH): + sr = Int64(cutlass.min(Int32(r0 + r), num_splits - 1)) + vs[r] = [ + _ldg_cg_f32x4(o_base + sr * 4096 + 16 * i) + for i in range(E // 4) + ] + for r in cutlass.range_constexpr(min(RCH, NSPLIT - r0)): + corr = Float32(0.0) + if r0 + r < num_splits: + corr = lds_f32(ml_smem + ((r0 + r) * 16 + 8 + eh) * 4) + for i in cutlass.range_constexpr(E // 4): + for j in cutlass.range_constexpr(4): + o[4 * i + j] = o[4 * i + j] + corr * vs[r][i][j] + inv = Float32(1.0) / (l_all + _ex2(sink_s - m_all_s)) + self.store_out(out_row, ed, o, inv, E) + + @cute.jit + def store_out(self, out_row, ed, o, inv, E: cutlass.Constexpr): + if cutlass.const_expr(E == 4): + _stg_u32x2( + out_row + ed * 2, + bf16x2(o[0] * inv, o[1] * inv), + bf16x2(o[2] * inv, o[3] * inv), + ) + for i in cutlass.range_constexpr(E // 8): + w = [ + bf16x2(o[8 * i + 2 * j] * inv, o[8 * i + 2 * j + 1] * inv) + for j in range(4) + ] + stg_u32x4(out_row + (ed + 8 * i) * 2, w[0], w[1], w[2], w[3]) + + @cute.jit + def issue(self, kv_head, skb, blk, dst, mbar, lane): + """Bulk-copy one 8 KB page (16 tokens x [K | V]) into a ring slot.""" + if lane == 0: + cute.arch.mbarrier_arrive_and_expect_tx(mbar, _PAGE_BYTES) + _bulk_copy( + dst, kv_head + Int64(blk) * skb * 2, smem_addr(mbar), _PAGE_BYTES + ) + + @cute.jit + def load_page(self, kv, slot, page, start_tok, seq_len, g, t): + """Lane fragments: K tokens g, g+8 (dims 32i + 8t ..) and V tokens 2t, + 2t+1, 2t+8, 2t+9 (dims 64h + 8g ..). V rows outside the attended range + are zeroed so stale cache contents never reach P V.""" + for r in cutlass.range_constexpr(2): + for i in cutlass.range_constexpr(4): + w = _lds128(slot + (g + 8 * r) * _ROW_BYTES + (32 * i + 8 * t) * 2) + for j in cutlass.range_constexpr(4): + kv[16 * r + 4 * i + j] = w[j] + vbase = slot + _HEAD_DIM * 2 + for r in cutlass.range_constexpr(4): + tok = 2 * t + (r % 2) + 8 * (r // 2) + for h in cutlass.range_constexpr(2): + w = _lds128(vbase + tok * _ROW_BYTES + (64 * h + 8 * g) * 2) + for j in cutlass.range_constexpr(4): + kv[32 + 4 * (2 * r + h) + j] = w[j] + p0 = page * _PAGE + if (p0 < start_tok) | (p0 + _PAGE > seq_len): + for r in cutlass.range_constexpr(4): + pos = p0 + 2 * t + (r % 2) + 8 * (r // 2) + ok = (pos >= start_tok) & (pos < seq_len) + for j in cutlass.range_constexpr(8): + kv[32 + 8 * r + j] = cutlass.select_( + ok, kv[32 + 8 * r + j], Uint32(0) + ) + + @cute.jit + def compute( + self, + kv, + qf, + acc, + page, + start_tok, + seq_len, + scale_log2, + m0, + m1, + l0, + l1, + g, + t, + lane, + ): + # S^T[tok, head] = K Q^T. Chunk c covers dims 32(c//2) + 8t + 4(c%2) + 0..3. + sa = [Float32(0.0)] * 4 + sb = [Float32(0.0)] * 4 + for c in cutlass.range_constexpr(8): + i = c // 2 + h = c % 2 + a = ( + kv[4 * i + 2 * h], + kv[16 + 4 * i + 2 * h], + kv[4 * i + 2 * h + 1], + kv[16 + 4 * i + 2 * h + 1], + ) + b0, b1 = qf[4 * i + 2 * h], qf[4 * i + 2 * h + 1] + if cutlass.const_expr(c % 2 == 0): + sa = mma(sa, a, b0, b1) + else: + sb = mma(sb, a, b0, b1) + s0 = sa[0] + sb[0] + s1 = sa[1] + sb[1] + s2 = sa[2] + sb[2] + s3 = sa[3] + sb[3] + pos = page * _PAGE + g + ninf = _neg_inf() + if (pos < start_tok) | (pos >= seq_len): + s0 = ninf + s1 = ninf + if (pos + 8 < start_tok) | (pos + 8 >= seq_len): + s2 = ninf + s3 = ninf + mx0 = cute.arch.fmax(s0, s2) + mx1 = cute.arch.fmax(s1, s3) + for off in cutlass.range_constexpr(3): + mx0 = cute.arch.fmax(mx0, cute.arch.shuffle_sync_bfly(mx0, 4 << off)) + mx1 = cute.arch.fmax(mx1, cute.arch.shuffle_sync_bfly(mx1, 4 << off)) + mn0 = cute.arch.fmax(m0, mx0) + mn1 = cute.arch.fmax(m1, mx1) + ms0 = mn0 * scale_log2 + ms1 = mn1 * scale_log2 + alpha0 = _ex2(fma(m0, scale_log2, -ms0)) + alpha1 = _ex2(fma(m1, scale_log2, -ms1)) + p0 = _ex2(fma(s0, scale_log2, -ms0)) + p1 = _ex2(fma(s1, scale_log2, -ms1)) + p2 = _ex2(fma(s2, scale_log2, -ms0)) + p3 = _ex2(fma(s3, scale_log2, -ms1)) + l0 = l0 * alpha0 + (p0 + p2) + l1 = l1 * alpha1 + (p1 + p3) + for j in cutlass.range_constexpr(8): + acc[4 * j] = acc[4 * j] * alpha0 + acc[4 * j + 1] = acc[4 * j + 1] * alpha1 + acc[4 * j + 2] = acc[4 * j + 2] * alpha0 + acc[4 * j + 3] = acc[4 * j + 3] * alpha1 + + # P^T as the B operand: lane (g, t) needs P[tok 2t, 2t+1, 2t+8, 2t+9][head g]. + x = bf16x2(p0, p1) + y = bf16x2(p2, p3) + src0 = 8 * t + (g >> 1) + x0 = cute.arch.shuffle_sync(x, src0) + x1 = cute.arch.shuffle_sync(x, src0 + 4) + y0 = cute.arch.shuffle_sync(y, src0) + y1 = cute.arch.shuffle_sync(y, src0 + 4) + sel = cutlass.select_((g & 1) == 0, Uint32(0x5410), Uint32(0x7632)) + b0 = _prmt(x0, x1, sel) + b1 = _prmt(y0, y1, sel) + + # O^T[dim, head] += V^T P^T. m-tile j rows g / g+8 are dims d0 / d0+1. + for j in cutlass.range_constexpr(8): + hh = j // 4 + wi = j % 4 + va = kv[32 + 4 * (0 * 2 + hh) + wi] + vb = kv[32 + 4 * (1 * 2 + hh) + wi] + vc = kv[32 + 4 * (2 * 2 + hh) + wi] + vd = kv[32 + 4 * (3 * 2 + hh) + wi] + acc[4 * j], acc[4 * j + 1], acc[4 * j + 2], acc[4 * j + 3] = mma( + (acc[4 * j], acc[4 * j + 1], acc[4 * j + 2], acc[4 * j + 3]), + ( + _prmt(va, vb, 0x5410), + _prmt(va, vb, 0x7632), + _prmt(vc, vd, 0x5410), + _prmt(vc, vd, 0x7632), + ), + b0, + b1, + ) + return mn0, mn1, l0, l1 + + +@cache +def _compile( + num_kv_heads, num_splits, num_warps, num_stages, sliding, min_pages, cluster +): + bf16 = cutlass.BFloat16 + batch = cute.sym_int() + q = make_fake_tensor( + bf16, (cute.sym_int(), num_kv_heads * _GQA, _HEAD_DIM), divisibility=8 + ) + kv = make_fake_tensor( + bf16, (cute.sym_int(), num_kv_heads, _PAGE, 2 * _HEAD_DIM), divisibility=8 + ) + bt = make_fake_tensor(Int32, (batch, cute.sym_int()), divisibility=1) + seq_lens = make_fake_tensor(Int32, (batch,), divisibility=1) + sinks = make_fake_tensor(Float32, (num_kv_heads * _GQA,), divisibility=1) + out = make_fake_tensor( + bf16, (cute.sym_int(), num_kv_heads * _GQA, _HEAD_DIM), divisibility=8 + ) + part_o = make_fake_tensor(Float32, (cute.sym_int(),), divisibility=4) + part_ml = make_fake_tensor(Float32, (cute.sym_int(),), divisibility=4) + counters = make_fake_tensor(Int32, (cute.sym_int(),), divisibility=1) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + kernel = _FlintDecodeAttention( + num_kv_heads, num_splits, num_warps, num_stages, sliding, min_pages, cluster + ) + return cute.compile( + kernel, + q, + kv, + bt, + seq_lens, + sinks, + out, + part_o, + part_ml, + counters, + Float32(1.0), + Int32(0), + stream, + options="--enable-tvm-ffi", + ) + + +# Upper bounds for the persistent workspace. Partials are only used with more +# than one (non-cluster) split, where num_bh * splits <= 2 * num_sms. +_MAX_PARTIALS_PER_SM = 4 +_MAX_BH = 1 << 16 + +_workspaces: dict = {} +_last_stream: dict = {} + + +def _workspace(device, num_partials, num_bh): + """Persistent split partials and arrival counters, allocated once per + device at a fixed maximum (never replaced, so captured CUDA graphs keep + valid pointers). Counters reset themselves and are zeroed only at + allocation; launches sharing them must be ordered on one stream, as within + a model forward; a launch from a different stream than the previous one + first waits for that stream.""" + if not torch.cuda.is_current_stream_capturing(): + stream = torch.cuda.current_stream(device) + last = _last_stream.get(device) + if last is not None and last != stream: + stream.wait_stream(last) + _last_stream[device] = stream + ws = _workspaces.get(device) + if ws is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Flint decode attention must run eagerly once before CUDA graph " + "capture." + ) + max_partials = _MAX_PARTIALS_PER_SM * _num_sms(device.index) + ws = ( + torch.empty( + max_partials * _GQA * _HEAD_DIM, dtype=torch.float32, device=device + ), + torch.empty(max_partials * 2 * _GQA, dtype=torch.float32, device=device), + torch.zeros(_MAX_BH, dtype=torch.int32, device=device), + ) + _workspaces[device] = ws + if ws[1].numel() < num_partials * 2 * _GQA or ws[2].numel() < num_bh: + raise ValueError( + f"Flint decode attention workspace too small for {num_partials} " + f"partials / {num_bh} (request, kv head) pairs" + ) + return ws + + +@cache +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def _launch_config(batch, num_kv_heads, sliding, num_sms): + """(num_splits, num_warps, num_stages, min_pages, cluster) for a static + batch size.""" + num_bh = batch * num_kv_heads + if sliding: + if batch <= 2: + return 8, 4, 2, 1, True + return (4, 4, 3, 1, True) if batch <= 16 else (1, 4, 3, 1, False) + if batch <= 2: + return max(1, num_sms // num_bh), 8, 2, 2, False + return max(1, 2 * num_sms // num_bh), 4, 2 if batch <= 4 else 3, 2, False + + +def _supports_kv_cache(kv_cache: torch.Tensor) -> bool: + """Whether ``kv_cache`` has the HND bf16 page layout the kernel assumes: + [blocks, kv_heads, 16, 2 * 128], contiguous pages, 16 B aligned.""" + return ( + kv_cache.dtype == torch.bfloat16 + and kv_cache.dim() == 4 + and kv_cache.shape[2] == _PAGE + and kv_cache.shape[3] == 2 * _HEAD_DIM + and kv_cache.stride(3) == 1 + and kv_cache.stride(2) == 2 * _HEAD_DIM + and kv_cache.stride(0) % 8 == 0 + and kv_cache.stride(1) % 8 == 0 + and kv_cache.data_ptr() % 16 == 0 + ) + + +def flint_decode_attention( + q: torch.Tensor, + kv_cache: torch.Tensor, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + sinks: torch.Tensor, + window: int | None, + scale: float, + out: torch.Tensor, +) -> torch.Tensor: + """Paged decode attention with sinks, one query token per request. + + CUDA-graph safe: the launch shape depends only on the batch size; sequence + lengths and block tables are read on the device. + + Args: + q: [B, num_heads, 128] bf16 queries. + kv_cache: [num_blocks, num_kv_heads, 16, 256] bf16 HND cache with K in + ``[..., :128]`` and V in ``[..., 128:]``. + block_tables: [B, max_blocks] int32. + seq_lens: [B] int32 KV lengths including the current token (0 writes + zeros). + sinks: [num_heads] fp32 sink logits. + window: Attend to the last ``window`` positions, or None for all. + scale: Softmax scale. + out: [B, num_heads, 128] bf16 output. + + Returns: + ``out``. + + """ + batch, num_heads, head_dim = q.shape + num_kv_heads = kv_cache.shape[1] + assert head_dim == _HEAD_DIM and num_heads == num_kv_heads * _GQA + assert _supports_kv_cache(kv_cache), ( + "Flint decode attention needs FlashInfer's HND bf16 KV cache with " + "16-token blocks" + ) + assert q.stride(-1) == 1 and out.stride(-1) == 1 + assert seq_lens.shape[0] == batch == out.shape[0] == block_tables.shape[0] + if batch == 0: + return out + sliding = window is not None + splits, num_warps, num_stages, min_pages, cluster = _launch_config( + batch, num_kv_heads, sliding, _num_sms(q.device.index) + ) + num_bh = batch * num_kv_heads + part_o, part_ml, counters = _workspace( + q.device, num_bh * splits if splits > 1 and not cluster else 0, num_bh + ) + kernel = _compile( + num_kv_heads, splits, num_warps, num_stages, sliding, min_pages, cluster + ) + kernel( + q, + kv_cache, + block_tables, + seq_lens, + sinks, + out, + part_o, + part_ml, + counters, + scale * _LOG2E, + window if sliding else 0, + ) + return out diff --git a/vllm/models/flint/ops/linear.py b/vllm/models/flint/ops/linear.py new file mode 100644 index 00000000000..25f4ac6dc04 --- /dev/null +++ b/vllm/models/flint/ops/linear.py @@ -0,0 +1,253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint dense linears with a small-M CuTe GEMV fast path. + +NVFP4 (W4A16): the GEMV (``nvfp4_gemv``) reads the very tensors the FlashInfer +CuTe-DSL kernel was prepared with (ModelOpt-packed weight, 128x4-swizzled +scales), so dispatching per token count costs no extra weight memory; larger +M stays on FlashInfer. BF16: the same GEMV kernel on BF16 weights where it +beats cuBLAS, cuBLAS otherwise. +""" + +from collections.abc import Iterable, Iterator +from typing import Any + +import torch +from torch import nn + +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear.nvfp4.flashinfer import ( + FlashInferCuteDslNvFp4W4A16LinearKernel, +) +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from . import nvfp4_gemv + +logger = init_logger(__name__) + + +def _swizzled_scale_buffer(weight_scale: torch.Tensor, n: int, k: int) -> torch.Tensor: + """The flat 128x4-swizzled scale buffer behind FlashInfer's strided view of + it (no copy).""" + nbytes = (n + 127) // 128 * 128 * ((k // 16 + 3) // 4 * 4) + storage = weight_scale.untyped_storage() + assert ( + weight_scale.dtype == torch.float8_e4m3fn + and weight_scale.storage_offset() == 0 + and storage.nbytes() == nbytes + ), "FlashInfer NVFP4 scales are not the flat 128x4-swizzled buffer" + return torch.empty(0, dtype=torch.uint8, device=weight_scale.device).set_( + storage, 0, (nbytes,), (1,) + ) + + +class FlintNvFp4LinearKernel(FlashInferCuteDslNvFp4W4A16LinearKernel): + """FlashInfer CuTe-DSL W4A16 GEMM with a small-M CuTe GEMV fast path.""" + + def process_weights_after_loading(self, layer: nn.Module) -> None: + super().process_weights_after_loading(layer) + n, k = layer.output_size_per_partition, layer.input_size_per_partition + assert ( + getattr(layer, "weights_padding_cols", 0) == 0 + and layer.weight.dtype == torch.uint8 + and tuple(layer.weight.shape) == (n, k // 2) + and layer.weight.is_contiguous() + ), "FlashInfer NVFP4 weight is not in the GEMV layout" + layer.flint_weight_scale = _swizzled_scale_buffer(layer.weight_scale, n, k) + # Host copy of the global scale for the large-M cuBLAS alpha. + layer.flint_global_scale = float(layer.weight_global_scale.item()) + layer.flint_bf16 = False + layer.flint_pdl = True + layer.flint_nk = (n, k) + _setup_gemv(layer, n, k, True, bf16=False) + + def apply_weights( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert bias is None + x_2d = x.reshape(-1, x.shape[-1]) + if x_2d.shape[0] >= _DEQUANT_MIN_M: + out = _dequant_gemm(layer, x_2d) + return out.view(*x.shape[:-1], out.shape[-1]) + if _use_gemv(layer, x_2d, silu_mul=False): + out = nvfp4_gemv.nvfp4_gemv( + x_2d, + layer.weight, + layer.flint_weight_scale, + layer.weight_global_scale, + ) + return out.view(*x.shape[:-1], out.shape[-1]) + return super().apply_weights(layer, x, bias) + + +class FlintBf16LinearMethod(UnquantizedLinearMethod): + """BF16 linear: CuTe GEMV for small M, cuBLAS otherwise.""" + + def __init__(self, pdl: bool = True) -> None: + super().__init__() + self.pdl = pdl + + def process_weights_after_loading(self, layer: nn.Module) -> None: + super().process_weights_after_loading(layer) + w = layer.weight + n, k = w.shape + assert w.dtype == torch.bfloat16 and w.is_contiguous() + layer.flint_bf16 = True + layer.flint_pdl = self.pdl + layer.flint_nk = (n, k) + _setup_gemv(layer, n, k, self.pdl, bf16=True) + + def apply( + self, + layer: nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert bias is None + x_2d = x.reshape(-1, x.shape[-1]) + if _use_gemv(layer, x_2d, silu_mul=False): + out = nvfp4_gemv.bf16_gemv(x_2d, layer.weight, pdl=self.pdl) + return out.view(*x.shape[:-1], out.shape[-1]) + return super().apply(layer, x, bias) + + +# Token counts from which exact BF16 dequantization + cuBLAS beats FlashInfer +# CuTe-DSL in the full model (GB200). At 8192 tokens it wins in isolation +# (o_proj 107 -> 86 us) but loses in-model (dequant ~9 us under contention). +_DEQUANT_MIN_M = 16384 + + +def _dequant_gemm(layer: nn.Module, x: torch.Tensor) -> torch.Tensor: + """``bf16(fp32(x @ bf16(e2m1 * e4m3).T) * global)`` via cuBLAS: the global + scale is the GEMM alpha, applied to the FP32 accumulator.""" + w = nvfp4_gemv.dequantize_rows(layer.weight, layer.flint_weight_scale) + out = torch.empty(x.shape[0], w.shape[0], dtype=torch.bfloat16, device=x.device) + matmul = torch.backends.cuda.matmul + prev = matmul.allow_bf16_reduced_precision_reduction + matmul.allow_bf16_reduced_precision_reduction = False # FP32 split-K sums + try: + torch.addmm(out, x, w.t(), beta=0, alpha=layer.flint_global_scale, out=out) + finally: + matmul.allow_bf16_reduced_precision_reduction = prev + return out + + +def _setup_gemv(layer: nn.Module, n: int, k: int, pdl: bool, bf16: bool) -> None: + """Enable (and compile) the GEMV where its tuned table lists the shape.""" + layer.flint_gemv = nvfp4_gemv.tuned(n, k, bf16=bf16) + layer.flint_gemv_silu_mul = layer.output_partition_sizes == [ + n // 2 + ] * 2 and nvfp4_gemv.tuned(n, k, silu_mul=True, bf16=bf16) + for silu_mul in (False, True): + if layer.flint_gemv_silu_mul if silu_mul else layer.flint_gemv: + nvfp4_gemv.warmup(n, k, silu_mul=silu_mul, pdl=pdl, bf16=bf16) + logger.info_once( + "Flint dense %s%s (N=%d, K=%d) -> Flint CuTe GEMV for small M", + "BF16" if bf16 else "W4A16", + " + SiLU-and-mul" if silu_mul else "", + n, + k, + ) + + +def _gemv_input_ok(x: torch.Tensor) -> bool: + # 16-byte vector loads of activation rows. + return x.stride(-1) == 1 and x.stride(0) % 8 == 0 and x.data_ptr() % 16 == 0 + + +def _use_gemv(layer: nn.Module, x: torch.Tensor, silu_mul: bool) -> bool: + """Whether the tuned dispatch table picks the GEMV for this M.""" + flag = layer.flint_gemv_silu_mul if silu_mul else layer.flint_gemv + if not flag or x.shape[0] == 0: + return False + n, k = layer.flint_nk + if nvfp4_gemv.pick_nt(n, k, x.shape[0], silu_mul, layer.flint_bf16) is None: + return False + assert x.dtype == torch.bfloat16 and _gemv_input_ok(x) + return True + + +def gate_up_silu_mul(layer: nn.Module, x: torch.Tensor) -> torch.Tensor | None: + """``silu_and_mul(layer(x))`` in one GEMV for small M, else None.""" + if not ( + getattr(layer, "flint_gemv_silu_mul", False) + and x.dim() == 2 + and _use_gemv(layer, x, silu_mul=True) + ): + return None + if layer.flint_bf16: + return nvfp4_gemv.bf16_gemv(x, layer.weight, silu_mul=True, pdl=layer.flint_pdl) + return nvfp4_gemv.nvfp4_gemv( + x, + layer.weight, + layer.flint_weight_scale, + layer.weight_global_scale, + silu_mul=True, + pdl=layer.flint_pdl, + ) + + +_DENSE_SUFFIXES = ("qkv_proj", "o_proj", "gate_up_proj", "down_proj") + + +# BF16 layers launched without PDL: the 4x larger weight prefetch of a BF16 +# o_proj slows the attention tail (+0.18 ms per d1 forward with PDL). +_BF16_NO_PDL_SUFFIXES = ("o_proj",) + + +def install_flint_dense_kernels(model: nn.Module) -> None: + """Route the dense linears of ``model`` through the Flint GEMV (launched + with PDL except the BF16 ``_BF16_NO_PDL_SUFFIXES``): NVFP4 linears get + ``FlintNvFp4LinearKernel``, unquantized BF16 projections + (``_DENSE_SUFFIXES``) ``FlintBf16LinearMethod``. Must run before weights + are processed after loading.""" + count = 0 + for module in model.modules(): + if not isinstance(module, LinearBase): + continue + method: Any = getattr(module, "quant_method", None) + kernel = getattr(method, "kernel", None) + if type(kernel) is FlashInferCuteDslNvFp4W4A16LinearKernel: + method.kernel = FlintNvFp4LinearKernel(kernel.config) + count += 1 + elif type(method) is UnquantizedLinearMethod and module.prefix.endswith( + _DENSE_SUFFIXES + ): + pdl = not module.prefix.endswith(_BF16_NO_PDL_SUFFIXES) + module.quant_method = FlintBf16LinearMethod(pdl) + count += 1 + logger.info("Flint dense GEMV installed on %d linear layers", count) + + +_MERGED_SHARDS = { + ".q_proj.": ".qkv_proj.", + ".k_proj.": ".qkv_proj.", + ".v_proj.": ".qkv_proj.", + ".gate_proj.": ".gate_up_proj.", + ".up_proj.": ".gate_up_proj.", +} + + +def check_merged_global_scales( + weights: Iterable[tuple[str, torch.Tensor]], +) -> Iterator[tuple[str, torch.Tensor]]: + """Pass checkpoint tensors through, raising if the NVFP4 global scales of + shards fused into one linear (q/k/v, gate/up) differ: the fused layer keeps + a single global scale.""" + seen: dict[str, torch.Tensor] = {} + for name, tensor in weights: + if name.endswith(".weight_scale_2"): + for shard, merged in _MERGED_SHARDS.items(): + if shard in name: + key = name.replace(shard, merged) + ref = seen.setdefault(key, tensor) + if not torch.equal(ref.reshape(()), tensor.reshape(())): + raise ValueError( + f"{name}: global scale {tensor.item()} differs from " + f"the other shards of {key} ({ref.item()})" + ) + break + yield name, tensor diff --git a/vllm/models/flint/ops/moe.py b/vllm/models/flint/ops/moe.py new file mode 100644 index 00000000000..bcf094ca85e --- /dev/null +++ b/vllm/models/flint/ops/moe.py @@ -0,0 +1,1127 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused NVFP4 W4A16 routed-expert MoE for Flint (plus decode kernels for the +BF16 checkpoint's TRTLLM-layout experts). + +Numerics: activations stay BF16; weights are dequantized exactly (E2M1 * E4M3) +and products are accumulated in fp32 with the per-expert fp32 global scale +applied to the fp32 accumulator. Router logits are BF16, top-k selection and +sigmoid/renormalization are fp32 (exact exp/division) and the routing weights +are rounded to BF16. Rounding points follow the reference (TRTLLM MoE rule): +gate/up stay fp32 and act = bf16(silu(g) * u), the sigmoid computed as +rcp.approx(1 + ex2.approx(-g * log2(e))) as in FlashInfer's kernel; each expert +output is y_e = bf16(acc_e * global_scale_e); the routed sum is +bf16(sum_e tw_e * y_e) accumulated in fp32; the output is bf16(routed + +shared). The decode, decode-block and grouped paths implement this rule. + +Decode (few tokens): three Triton kernels, (1) router GEMV + top-k, (2) per +(token, expert) gate/up GEMV + SiLU * up, (3) down GEMV over the top-k experts ++ weighted combine + shared-expert add. Larger batches: cuBLAS router GEMM + +Triton top-k, FlashInfer's SM100 W4A16 grouped GEMMs (see ``moe_cute``) and a +Triton combine. Both paths read the same weights in FlashInfer's layout (see +``FlintMoEWeights.from_modelopt``). +""" + +from collections.abc import Iterable +from dataclasses import dataclass + +import regex as re +import torch +from torch import nn + +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + interleave_linear_and_gate, + nvfp4_swizzled_scale_to_cutedsl_mma_view, + reorder_w13_to_w31_for_flashinfer_cutedsl, +) +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( + swizzle_blockscale, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, tldevice, triton +from vllm.utils.torch_utils import aux_stream, current_stream + +from . import moe_cute, moe_gemv +from .norm import _rsqrt + +# Largest token count served by the GEMV decode kernels. +DECODE_MAX_TOKENS = 6 +# Largest token count served by the BF16-expert decode kernels. +BF16_DECODE_MAX_TOKENS = 12 +# Largest token count routed by the router GEMV kernel (cuBLAS above). +ROUTE_GEMV_MAX_TOKENS = 64 + + +def _ptxas_has_bf16x2_e2m1() -> bool: + # cvt.rn.bf16x2.e2m1x2 needs PTX ISA 9.2 (ptxas >= 13.2); some images run + # Triton with an older system ptxas. + if not current_platform.is_cuda(): + return False + from triton.backends.nvidia.compiler import get_ptxas + + major, minor = map(int, get_ptxas(100).version.split(".")[:2]) + return (major, minor) >= (13, 2) + + +# Eight E2M1 codes of one int32 word -> eight bf16 or fp16 (code value * +# scale): the hardware E2M1x2 conversion then a multiply by the E4M3 block +# scale. Both are exact in either type: E2M1 * E4M3 has at most 6 significant +# bits and lies in [2^-10, 2688]. BF16 is faster (its fp32 widening is free), +# FP16 assembles with older ptxas. +_DEQUANT_BF16 = tl.constexpr(_ptxas_has_bf16x2_e2m1()) +_DEQUANT_T = "bf16" if _DEQUANT_BF16 else "f16" +_DEQUANT_ASM = tl.constexpr( + f"""{{ +.reg .b8 b0, b1, b2, b3; +.reg .b32 p0, p1, p2, p3; +mov.b32 {{b0, b1, b2, b3}}, $8; +cvt.rn.{_DEQUANT_T}x2.e2m1x2 p0, b0; +cvt.rn.{_DEQUANT_T}x2.e2m1x2 p1, b1; +cvt.rn.{_DEQUANT_T}x2.e2m1x2 p2, b2; +cvt.rn.{_DEQUANT_T}x2.e2m1x2 p3, b3; +mul.rn.{_DEQUANT_T}x2 p0, p0, $9; +mul.rn.{_DEQUANT_T}x2 p1, p1, $9; +mul.rn.{_DEQUANT_T}x2 p2, p2, $9; +mul.rn.{_DEQUANT_T}x2 p3, p3, $9; +mov.b32 {{$0, $1}}, p0; +mov.b32 {{$2, $3}}, p1; +mov.b32 {{$4, $5}}, p2; +mov.b32 {{$6, $7}}, p3; +}}""" +) + + +@triton.jit +def _silu(g): + # g * sigmoid(g), sigmoid = rcp.approx(1 + ex2.approx(-g * log2(e))): the + # instructions of FlashInfer's W4A16 MoE epilogue (and TRTLLM's). + e = tl.inline_asm_elementwise( + "ex2.approx.ftz.f32 $0, $1;", + "=f,f", + [g * -1.4426950408889634], + tl.float32, + True, + 1, + ) + sig = tl.inline_asm_elementwise( + "rcp.approx.ftz.f32 $0, $1;", "=f,f", [1.0 + e], tl.float32, True, 1 + ) + return g * sig + + +@triton.jit +def _dequant8(q, s): + """int32 words of E2M1 codes + per-word E4M3 scales (uint8), same shape -> + eight bf16 or fp16 tensors (code j of each word, value * scale, exact).""" + dt = tl.bfloat16 if _DEQUANT_BF16 else tl.float16 + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32).to(dt) + s = s.to(tl.uint16, bitcast=True).to(tl.uint32) + s = (s | (s << 16)).to(tl.int32, bitcast=True) + return tl.inline_asm_elementwise( + _DEQUANT_ASM, + "=h,=h,=h,=h,=h,=h,=h,=h,r,r", + [q, s], + dtype=(dt,) * 8, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _dequant_rows(q, s, BN: tl.constexpr, BK: tl.constexpr): + """[BN, BK / 8] int32 code words + per-word E4M3 scales -> [BN, BK] bf16 + or fp16 (value * scale, exact), K in natural order.""" + o = _dequant8(q, s) + w = tl.join( + tl.join(tl.join(o[0], o[4]), tl.join(o[2], o[6])), + tl.join(tl.join(o[1], o[5]), tl.join(o[3], o[7])), + ) + return tl.reshape(w, [BN, BK]) + + +@triton.jit +def _scale_offsets(e, rows, words, ROWS: tl.constexpr, K: tl.constexpr): + """Byte offsets of the E4M3 block scales of (rows, int32 words) in + FlashInfer's MMA scale layout: physical (E, ROWS / 128, K / 64, 32, 4, 4) + for logical (row % 32, row % 128 // 32, row // 128, group % 4, group // 4).""" + g = words // 2 + return ( + ((e * (ROWS // 128) + rows[:, None] // 128) * (K // 64) + g[None, :] // 4) * 512 + + (rows[:, None] % 32) * 16 + + (rows[:, None] % 128) // 32 * 4 + + g[None, :] % 4 + ) + + +@triton.jit +def _topk_row( + logits, + offs, + ids_ptr, + tw_ptr, + E: tl.constexpr, + TOPK: tl.constexpr, + TOPK_P2: tl.constexpr, +): + """Top-k of one row of bf16 logits (ties -> lowest expert index), fp32 + sigmoid scores renormalized over the selected experts, ``s / (sum + + 1e-20)``, rounded to bf16 (the reference packs bf16 routing weights). + + Each logit and its index are packed into one order-preserving int32 key so + every selection step is a single max-reduction. + """ + tl.static_assert(E <= 256) + bits = logits.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF + ordered = tl.where(bits >= 0x8000, bits ^ 0xFFFF, bits | 0x8000) + keys = (ordered << 8) | (E - 1 - offs) + koffs = tl.arange(0, TOPK_P2) + sel_ids = tl.zeros([TOPK_P2], dtype=tl.int32) + sel_s = tl.zeros([TOPK_P2], dtype=tl.float32) + for i in tl.static_range(TOPK): + best = tl.max(keys, axis=0) + o = best >> 8 + b16 = tl.where(o >= 0x8000, o ^ 0x8000, o ^ 0xFFFF) + logit = b16.to(tl.int16).to(tl.bfloat16, bitcast=True).to(tl.float32) + s = tl.math.div_rn(1.0, 1.0 + tldevice.exp(-logit)) + sel_ids = tl.where(koffs == i, E - 1 - (best & 0xFF), sel_ids) + sel_s = tl.where(koffs == i, s, sel_s) + keys = tl.where(keys == best, -1, keys) + kmask = koffs < TOPK + tl.store(ids_ptr + koffs, sel_ids, mask=kmask) + tw = tl.math.div_rn(sel_s, tl.sum(sel_s, axis=0) + 1e-20) + tl.store(tw_ptr + koffs, tw.to(tl.bfloat16).to(tl.float32), mask=kmask) + + +@triton.jit +def _add_rms_norm(a_ptr, r_ptr, w_ptr, t, eps, H: tl.constexpr): + """``norm.fused_add_rms_norm`` for row t: (bf16 residual, bf16 out), both + from the unrounded fp32 sum s = a + r.""" + offs = tl.arange(0, H) + a = tl.load(a_ptr + t * H + offs).to(tl.float32) + r = tl.load(r_ptr + t * H + offs).to(tl.float32) + s = a + r + inv = _rsqrt(tl.sum(s * s, axis=0) / H + eps) + w = tl.load(w_ptr + offs).to(tl.float32) + return s.to(tl.bfloat16), (s * inv * w).to(tl.bfloat16) + + +@triton.jit +def _router_kernel( + x_ptr, + res_ptr, + ln_w_ptr, + eps, + res_out_ptr, + xn_ptr, + wr_ptr, + logits_ptr, + ids_ptr, + tw_ptr, + counter_ptr, + H: tl.constexpr, + E: tl.constexpr, + TOPK: tl.constexpr, + EB: tl.constexpr, + FUSE_NORM: tl.constexpr, +): + """Router GEMV (bf16 logits) of experts [pe * EB, + EB) for token t, then + top-k by the token's last program. Grid (M, E // EB), or (M * E // EB,) + with FUSE_NORM, where the GEMV input is ``norm.fused_add_rms_norm`` of x and + res (program pe == 0 writes the residual and normed rows). Router weights + are fetched before the PDL dependency wait. + + The last program is found with a per-token atomic counter that is never + reset (the low bits of the signed int32 count; wraparound-safe since the + program count is a power of two), so the buffer must be zeroed once at + allocation and only used by launches serialized on one stream. + """ + NR: tl.constexpr = E // EB + tl.static_assert((NR & (NR - 1)) == 0) + if FUSE_NORM: + t = tl.program_id(0) // NR + pe = tl.program_id(0) % NR + else: + t = tl.program_id(0) + pe = tl.program_id(1) + offs_e = pe * EB + tl.arange(0, EB) + offs_k = tl.arange(0, H) + w = tl.load(wr_ptr + offs_e[:, None] * H + offs_k[None, :]).to(tl.float32) + tl.extra.cuda.gdc_wait() + if FUSE_NORM: + s, xn = _add_rms_norm(x_ptr, res_ptr, ln_w_ptr, t, eps, H) + if pe == 0: + tl.store(res_out_ptr + t * H + offs_k, s) + tl.store(xn_ptr + t * H + offs_k, xn) + x = xn.to(tl.float32) + else: + x = tl.load(x_ptr + t * H + offs_k).to(tl.float32) + logits = tl.sum(w * x[None, :], axis=1).to(tl.bfloat16) + tl.store(logits_ptr + t * E + offs_e, logits) + tl.debug_barrier() + count = tl.atomic_add(counter_ptr + t, 1, sem="acq_rel") + if (count & (NR - 1)) == NR - 1: + offs = tl.arange(0, E) + vals = tl.load(logits_ptr + t * E + offs, volatile=True) + _topk_row( + vals, + offs, + ids_ptr + t * TOPK, + tw_ptr + t * TOPK, + E, + TOPK, + triton.next_power_of_2(TOPK), + ) + + +@triton.jit +def _topk_kernel( + logits_ptr, + ids_ptr, + tw_ptr, + E: tl.constexpr, + TOPK: tl.constexpr, +): + """Top-k of bf16 router logits. Grid (M,).""" + m = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, E) + vals = tl.load(logits_ptr + m * E + offs) + _topk_row( + vals, + offs, + ids_ptr + m * TOPK, + tw_ptr + m * TOPK, + E, + TOPK, + triton.next_power_of_2(TOPK), + ) + + +@triton.jit +def _gate_up_gemv_kernel( + x_ptr, + ids_ptr, + w_ptr, + s_ptr, + gs_ptr, + act_ptr, + H: tl.constexpr, + INTER: tl.constexpr, + TOPK: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + """act[p] = silu(x[t] @ Wg[e]^T) * (x[t] @ Wu[e]^T), pair p = t * TOPK + j. + + w_ptr: int32 view of the FlashInfer GEMM1 weight ([up; gate] interleaved + in 64-row groups). Grid (M * TOPK, INTER // BN). + """ + tl.extra.cuda.gdc_wait() + p = tl.program_id(0) + pn = tl.program_id(1) + t = p // TOPK + e = tl.load(ids_ptr + p).to(tl.int64) + cols = pn * BN + tl.arange(0, BN) + rows_u = cols // 64 * 128 + cols % 64 + rows_g = rows_u + 64 + offs_w = tl.arange(0, BK // 8) + w_base = w_ptr + e * (2 * INTER * H // 8) + offs_w[None, :] + # Full-shape (broadcast) activation pointers keep x in the weight tile's + # layout (vectorized loads, no layout conversion); repeats hit L1. + x_ptrs = x_ptr + t * H + tl.arange(0, BK)[None, :] + x_ptrs += tl.zeros([BN, BK], dtype=tl.int32) + acc_g = tl.zeros([BN, BK], dtype=tl.float32) + acc_u = tl.zeros([BN, BK], dtype=tl.float32) + for k0 in tl.static_range(0, H, BK): + words = k0 // 8 + offs_w + qg = tl.load(w_base + rows_g[:, None] * (H // 8) + k0 // 8) + qu = tl.load(w_base + rows_u[:, None] * (H // 8) + k0 // 8) + sg = tl.load(s_ptr + _scale_offsets(e, rows_g, words, 2 * INTER, H)) + su = tl.load(s_ptr + _scale_offsets(e, rows_u, words, 2 * INTER, H)) + xk = tl.load(x_ptrs + k0).to(tl.float32) + acc_g += _dequant_rows(qg, sg, BN, BK).to(tl.float32) * xk + acc_u += _dequant_rows(qu, su, BN, BK).to(tl.float32) * xk + gs = tl.load(gs_ptr + e) + g = tl.sum(acc_g, axis=1) * gs + u = tl.sum(acc_u, axis=1) * gs + tl.store(act_ptr + p * INTER + cols, (_silu(g) * u).to(tl.bfloat16)) + + +@triton.jit +def _down_gemv_kernel( + act_ptr, + ids_ptr, + tw_ptr, + w_ptr, + s_ptr, + gs_ptr, + shared_ptr, + out_ptr, + H: tl.constexpr, + INTER: tl.constexpr, + TOPK: tl.constexpr, + BN: tl.constexpr, +): + """out[t] = bf16(bf16(sum_j tw[t, j] * y_j) + shared[t]) with + y_j = bf16(act[t * TOPK + j] @ W2[e_j]^T), weighted sum in fp32. + + w_ptr: int32 view of the down weight. Grid (M, H // BN). + """ + tl.extra.cuda.gdc_wait() + t = tl.program_id(0) + pn = tl.program_id(1) + rows = pn * BN + tl.arange(0, BN) + offs_w = tl.arange(0, INTER // 8) + out = tl.zeros([BN], dtype=tl.float32) + for j in tl.static_range(TOPK): + p = t * TOPK + j + e = tl.load(ids_ptr + p).to(tl.int64) + q = tl.load( + w_ptr + + e * (H * INTER // 8) + + rows[:, None] * (INTER // 8) + + offs_w[None, :] + ) + sc = tl.load(s_ptr + _scale_offsets(e, rows, offs_w, H, INTER)) + a_ptrs = act_ptr + p * INTER + tl.arange(0, INTER)[None, :] + a = tl.load(a_ptrs + tl.zeros([BN, INTER], dtype=tl.int32)).to(tl.float32) + acc = _dequant_rows(q, sc, BN, INTER).to(tl.float32) * a + gs = tl.load(gs_ptr + e) + y = (tl.sum(acc, axis=1) * gs).to(tl.bfloat16).to(tl.float32) + out += tl.load(tw_ptr + p) * y + shared = tl.load(shared_ptr + t * H + rows).to(tl.float32) + out = (out.to(tl.bfloat16).to(tl.float32) + shared).to(tl.bfloat16) + tl.store(out_ptr + t * H + rows, out) + + +@triton.jit +def _combine_kernel( + y_ptr, + perm_ptr, + tw_ptr, + shared_ptr, + out_ptr, + H: tl.constexpr, + TOPK: tl.constexpr, + BH: tl.constexpr, +): + """out[t] = bf16(bf16(sum_j tw[t, j] * y[perm[t * TOPK + j]]) + shared[t]), + weighted sum in fp32. + + Grid (M, H // BH). + """ + tl.extra.cuda.gdc_wait() + t = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BH + tl.arange(0, BH) + acc = tl.zeros([BH], dtype=tl.float32) + for j in tl.static_range(TOPK): + row = tl.load(perm_ptr + t * TOPK + j).to(tl.int64) + y = tl.load(y_ptr + row * H + offs).to(tl.float32) + acc += tl.load(tw_ptr + t * TOPK + j) * y + shared = tl.load(shared_ptr + t * H + offs).to(tl.float32) + out = (acc.to(tl.bfloat16).to(tl.float32) + shared).to(tl.bfloat16) + tl.store(out_ptr + t * H + offs, out) + + +@triton.jit +def _block_major_offsets(rows, k, ROWS: tl.constexpr): + """Element offsets of (rows, k) in FlashInfer's BF16 BlockMajorK expert + layout: [K / 64, ROWS, 64].""" + return (k[None, :] // 64) * (ROWS * 64) + rows[:, None] * 64 + k[None, :] % 64 + + +@triton.jit +def _gate_up_gemv_bf16_kernel( + x_ptr, + ids_ptr, + w_ptr, + rowmap_ptr, + act_ptr, + H: tl.constexpr, + INTER: tl.constexpr, + TOPK: tl.constexpr, + BN: tl.constexpr, + BK: tl.constexpr, +): + """BF16 experts: act[p] = silu(x[t] @ Wg[e]^T) * (x[t] @ Wu[e]^T). + + w_ptr: TRTLLM BlockMajorK gate/up weight [E, H / 64, 2I, 64]; rowmap: + logical row (gate rows, then up rows) -> physical row. Grid + (M * TOPK, INTER // BN). + """ + tl.extra.cuda.gdc_wait() + p = tl.program_id(0) + pn = tl.program_id(1) + t = p // TOPK + e = tl.load(ids_ptr + p).to(tl.int64) + cols = pn * BN + tl.arange(0, BN) + rows_g = tl.load(rowmap_ptr + cols) + rows_u = tl.load(rowmap_ptr + INTER + cols) + offs_k = tl.arange(0, BK) + w_base = w_ptr + e * (2 * INTER * H) + x_ptrs = x_ptr + t * H + offs_k[None, :] + tl.zeros([BN, BK], dtype=tl.int32) + acc_g = tl.zeros([BN, BK], dtype=tl.float32) + acc_u = tl.zeros([BN, BK], dtype=tl.float32) + for k0 in tl.static_range(0, H, BK): + k = k0 + offs_k + wg = tl.load(w_base + _block_major_offsets(rows_g, k, 2 * INTER)) + wu = tl.load(w_base + _block_major_offsets(rows_u, k, 2 * INTER)) + xk = tl.load(x_ptrs + k0).to(tl.float32) + acc_g += wg.to(tl.float32) * xk + acc_u += wu.to(tl.float32) * xk + g = tl.sum(acc_g, axis=1) + u = tl.sum(acc_u, axis=1) + # TRTLLM BF16 MoE semantics (the BF16 baseline used for larger batches): + # gate/up stay fp32, one bf16 rounding of silu(g) * u. + tl.store(act_ptr + p * INTER + cols, (_silu(g) * u).to(tl.bfloat16)) + + +@triton.jit +def _down_gemv_bf16_kernel( + act_ptr, + ids_ptr, + tw_ptr, + w_ptr, + rowmap_ptr, + shared_ptr, + out_ptr, + H: tl.constexpr, + INTER: tl.constexpr, + TOPK: tl.constexpr, + BN: tl.constexpr, +): + """BF16 experts: out[t] = bf16(bf16(sum_j tw[t, j] * y_j) + shared[t]), + y_j = bf16(act[t * TOPK + j] @ W2[e_j]^T). + + w_ptr: TRTLLM BlockMajorK down weight [E, I / 64, H, 64]; rowmap: output + row -> physical row. Grid (M, H // BN). + """ + tl.extra.cuda.gdc_wait() + t = tl.program_id(0) + pn = tl.program_id(1) + rows = pn * BN + tl.arange(0, BN) + prow = tl.load(rowmap_ptr + rows) + offs_k = tl.arange(0, INTER) + w_offs = _block_major_offsets(prow, offs_k, H) + out = tl.zeros([BN], dtype=tl.float32) + for j in tl.static_range(TOPK): + p = t * TOPK + j + e = tl.load(ids_ptr + p).to(tl.int64) + w = tl.load(w_ptr + e * (H * INTER) + w_offs).to(tl.float32) + a_ptrs = act_ptr + p * INTER + offs_k[None, :] + a = tl.load(a_ptrs + tl.zeros([BN, INTER], dtype=tl.int32)).to(tl.float32) + y = tl.sum(w * a, axis=1).to(tl.bfloat16).to(tl.float32) + out += tl.load(tw_ptr + p) * y + shared = tl.load(shared_ptr + t * H + rows).to(tl.float32) + out = (out.to(tl.bfloat16).to(tl.float32) + shared).to(tl.bfloat16) + tl.store(out_ptr + t * H + rows, out) + + +def trtllm_bf16_row_maps( + hidden_size: int, intermediate_size: int, device: torch.device | str +) -> tuple[torch.Tensor, torch.Tensor]: + """Logical -> physical row maps of the TRTLLM BF16 expert weights as + produced by ``convert_moe_weights_to_flashinfer_trtllm_block_layout`` + (gated gate/up interleave + epilogue-tile shuffle). Returns (w13 map over + [gate rows; up rows], w2 map over output rows), int32.""" + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + + rows13 = 2 * intermediate_size + w13 = torch.empty(rows13, hidden_size * 2, dtype=torch.uint8, device=device) + w2 = torch.empty( + hidden_size, intermediate_size * 2, dtype=torch.uint8, device=device + ) + src13 = _maybe_get_cached_w3_w1_permute_indices( + {}, w13, 128, is_gated_act_gemm=True + ) + src13 = (src13.to(device) + rows13 // 2) % rows13 + src2 = get_w2_permute_indices_with_cache({}, w2, 128).to(device) + # Physical row j holds logical row src[j]; invert. + return ( + torch.argsort(src13).to(torch.int32).contiguous(), + torch.argsort(src2).to(torch.int32).contiguous(), + ) + + +def flint_moe_bf16_decode( + x: torch.Tensor, + router_weight: torch.Tensor, + w13: torch.Tensor, + w2: torch.Tensor, + row_maps: tuple[torch.Tensor, torch.Tensor], + top_k: int, + counter: torch.Tensor, + shared_output: torch.Tensor, + shared_ready: torch.cuda.Event | None = None, +) -> torch.Tensor: + """Decode-path MoE on TRTLLM-layout BF16 experts (w13 [E, 2I, H], w2 + [E, H, I] storage in BlockMajorK order) with the TRTLLM BF16 MoE rounding + points: act = bf16(silu(g) * u) from fp32 gate/up, y_e = bf16(acc), + routed = bf16(sum tw * y_e) in fp32, out = bf16(routed + shared).""" + num_tokens, hidden = x.shape + inter = w2.shape[2] + assert x.dtype == torch.bfloat16 and x.is_contiguous() + ids, tw = route(x, router_weight, top_k, counter) + act = torch.empty(num_tokens * top_k, inter, dtype=x.dtype, device=x.device) + _gate_up_gemv_bf16_kernel[(num_tokens * top_k, inter)]( + x, + ids, + w13, + row_maps[0], + act, + H=hidden, + INTER=inter, + TOPK=top_k, + BN=1, + BK=256, + num_warps=1, + launch_pdl=True, + ) + out = torch.empty_like(x) + if shared_ready is not None: + shared_ready.wait(current_stream()) + bn = 1 if num_tokens <= 2 else 2 + _down_gemv_bf16_kernel[(num_tokens, hidden // bn)]( + act, + ids, + tw, + w2, + row_maps[1], + shared_output, + out, + H=hidden, + INTER=inter, + TOPK=top_k, + BN=bn, + num_warps=1, + launch_pdl=True, + ) + return out + + +@dataclass +class FlintMoEWeights: + """Routed-expert NVFP4 weights in FlashInfer's W4A16 layout.""" + + w1: torch.Tensor # uint8 [E, 2I, H/2], [up; gate] interleaved by 64 rows + w1_sf: torch.Tensor # float8_e4m3fn, 6D MMA scale layout + w1_alpha: torch.Tensor # fp32 [E] global scale of gate/up + w2: torch.Tensor # uint8 [E, H, I/2] + w2_sf: torch.Tensor + w2_alpha: torch.Tensor + + @classmethod + def from_modelopt( + cls, + w13: torch.Tensor, + w13_scale: torch.Tensor, + w13_scale_2: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + w2_scale_2: torch.Tensor, + ) -> "FlintMoEWeights": + """From ModelOpt tensors: w13 [E, 2I, H/2] ([gate; up]), w13_scale + [E, 2I, H/16], w2 [E, H, I/2], w2_scale [E, H, I/16], fp32 [E].""" + w1, w1_sf = reorder_w13_to_w31_for_flashinfer_cutedsl( + MoEActivation.SILU, w13, w13_scale + ) + w1 = interleave_linear_and_gate(w1, group_size=64, dim=1) + w1_sf = interleave_linear_and_gate(w1_sf, group_size=64, dim=1) + return cls( + w1, + nvfp4_swizzled_scale_to_cutedsl_mma_view(swizzle_blockscale(w1_sf)), + w13_scale_2.float().contiguous(), + w2.contiguous(), + nvfp4_swizzled_scale_to_cutedsl_mma_view(swizzle_blockscale(w2_scale)), + w2_scale_2.float().contiguous(), + ) + + @property + def intermediate_size(self) -> int: + return self.w2.shape[2] * 2 + + +def make_route_counter(device: torch.device | str | None = None) -> torch.Tensor: + """Zeroed counter buffer for the router GEMV kernel; allocate it outside + CUDA graph capture and use it from one stream only.""" + return torch.zeros(ROUTE_GEMV_MAX_TOKENS, dtype=torch.int32, device=device) + + +def route( + x: torch.Tensor, + router_weight: torch.Tensor, + top_k: int, + counter: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """BF16 router logits -> (top-k ids int32, renormalized fp32 weights).""" + num_tokens, hidden = x.shape + num_experts = router_weight.shape[0] + ids = torch.empty(num_tokens, top_k, dtype=torch.int32, device=x.device) + tw = torch.empty(num_tokens, top_k, dtype=torch.float32, device=x.device) + if num_tokens <= ROUTE_GEMV_MAX_TOKENS: + logits = torch.empty( + num_tokens, num_experts, dtype=torch.bfloat16, device=x.device + ) + eb = 2 + _router_kernel[(num_tokens, num_experts // eb)]( + x, + x, + x, + 0.0, + x, + x, + router_weight, + logits, + ids, + tw, + counter, + H=hidden, + E=num_experts, + TOPK=top_k, + EB=eb, + FUSE_NORM=False, + num_warps=4, + launch_pdl=True, + ) + else: + logits = torch.mm(x, router_weight.t()) + _topk_kernel[(num_tokens,)]( + logits, ids, tw, E=num_experts, TOPK=top_k, num_warps=4 + ) + return ids, tw + + +def flint_moe( + x: torch.Tensor, + router_weight: torch.Tensor, + weights: FlintMoEWeights, + top_k: int, + counter: torch.Tensor, + shared_output: torch.Tensor, + shared_ready: torch.cuda.Event | None = None, +) -> torch.Tensor: + """Routed-expert MoE output + shared_output for BF16 x [M, H]. + + counter: buffer from ``make_route_counter``. shared_ready: event the + current stream waits on before the first read of shared_output (which + may be produced on another stream). + """ + num_tokens, hidden = x.shape + inter = weights.intermediate_size + assert x.dtype == torch.bfloat16 and x.is_contiguous() + ids, tw = route(x, router_weight, top_k, counter) + out = torch.empty_like(x) + if num_tokens <= DECODE_MAX_TOKENS: + act = torch.empty(num_tokens * top_k, inter, dtype=x.dtype, device=x.device) + _gate_up_gemv_kernel[(num_tokens * top_k, inter)]( + x, + ids, + weights.w1.view(torch.int32), + weights.w1_sf, + weights.w1_alpha, + act, + H=hidden, + INTER=inter, + TOPK=top_k, + BN=1, + BK=256, + num_warps=1, + launch_pdl=True, + ) + bn = 1 if num_tokens <= 2 else 2 + if shared_ready is not None: + shared_ready.wait(current_stream()) + _down_gemv_kernel[(num_tokens, hidden // bn)]( + act, + ids, + tw, + weights.w2.view(torch.int32), + weights.w2_sf, + weights.w2_alpha, + shared_output, + out, + H=hidden, + INTER=inter, + TOPK=top_k, + BN=bn, + num_warps=1, + launch_pdl=True, + ) + return out + y, perm = moe_cute.grouped_moe( + x, + ids, + tw, + weights.w1, + weights.w1_sf, + weights.w1_alpha, + weights.w2, + weights.w2_sf, + weights.w2_alpha, + ) + if shared_ready is not None: + shared_ready.wait(current_stream()) + bh = 1024 + _combine_kernel[(num_tokens, hidden // bh)]( + y, + perm, + tw, + shared_output, + out, + H=hidden, + TOPK=top_k, + BH=bh, + num_warps=4, + launch_pdl=True, + ) + return out + + +_EXPERT_PARAM = re.compile( + r"(\d+)\.(gate_proj|up_proj|down_proj)\.(weight|weight_scale|weight_scale_2)" +) + + +def _start_shared( + shared: torch.Tensor | nn.Module, + x: torch.Tensor, + fork: torch.cuda.Event, + join: torch.cuda.Event, +) -> tuple[torch.Tensor, torch.cuda.Event | None]: + """Run a shared-expert module on the auxiliary stream (fork/join events, + CUDA-graph safe). Returns (output, event to wait on before reading it).""" + if not isinstance(shared, nn.Module): + return shared, None + stream = aux_stream() + fork.record(current_stream()) + with torch.cuda.stream(stream): + fork.wait(stream) + out = shared(x) + join.record(stream) + return out, join + + +class FlintBf16DecodeMoE(nn.Module): + """Decode path (few tokens) for Flint's BF16 routed experts. + + Reads the expert weights of a vLLM FusedMoE layer that uses the FlashInfer + TRTLLM BF16 backend (BlockMajorK, rows shuffled), so no extra weight copy + is needed; larger batches keep using that layer. + """ + + def __init__( + self, num_experts: int, hidden_size: int, intermediate_size: int, top_k: int + ) -> None: + super().__init__() + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.top_k = top_k + device = torch.empty(0).device + self.register_buffer("route_counter", make_route_counter(device), False) + map13, map2 = trtllm_bf16_row_maps(hidden_size, intermediate_size, device) + self.register_buffer("row_map13", map13, persistent=False) + self.register_buffer("row_map2", map2, persistent=False) + self._fork = torch.cuda.Event() + self._join = torch.cuda.Event() + self._checked = False + + def _check(self, routed_experts: nn.Module) -> None: + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + ) + + backend = getattr(routed_experts.quant_method, "unquantized_backend", None) + e, h, i = self.num_experts, self.hidden_size, self.intermediate_size + if ( + backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM + or tuple(routed_experts.w13_weight.shape) != (e, 2 * i, h) + or tuple(routed_experts.w2_weight.shape) != (e, h, i) + ): + raise RuntimeError( + "Flint BF16 decode MoE requires the FlashInfer TRTLLM BF16 MoE " + f"backend without padding (got {backend})" + ) + self._checked = True + + def forward( + self, + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + routed_experts: nn.Module, + shared: torch.Tensor | nn.Module, + ) -> torch.Tensor: + if not self._checked: + self._check(routed_experts) + shared, ready = _start_shared(shared, hidden_states, self._fork, self._join) + return flint_moe_bf16_decode( + hidden_states, + router_weight, + routed_experts.w13_weight, + routed_experts.w2_weight, + (self.row_map13, self.row_map2), + self.top_k, + self.route_counter, + shared, + ready, + ) + + +class FlintNvfp4Experts(nn.Module): + """Routed experts of one Flint MoE layer (NVFP4 GGUF checkpoints). + + Loads per-expert ModelOpt NVFP4 tensors (``{e}.{gate,up,down}_proj. + {weight, weight_scale, weight_scale_2}``) and, once every expert is loaded, + repacks them into ``FlintMoEWeights`` and frees the checkpoint layout. + """ + + def __init__( + self, + num_experts: int, + hidden_size: int, + intermediate_size: int, + top_k: int, + max_num_tokens: int = 0, + ) -> None: + """max_num_tokens: largest batch; the grouped-GEMM kernels for every + tactic up to it are compiled when the weights are loaded.""" + super().__init__() + self.max_num_tokens = max_num_tokens + e, h, i = num_experts, hidden_size, intermediate_size + self.num_experts = e + self.top_k = top_k + + def param(*shape: int, dtype: torch.dtype) -> nn.Parameter: + return nn.Parameter(torch.empty(*shape, dtype=dtype), requires_grad=False) + + self.w13_weight = param(e, 2 * i, h // 2, dtype=torch.uint8) + self.w13_weight_scale = param(e, 2 * i, h // 16, dtype=torch.float8_e4m3fn) + self.w13_weight_scale_2 = param(e, 2, dtype=torch.float32) + self.w2_weight = param(e, h, i // 2, dtype=torch.uint8) + self.w2_weight_scale = param(e, h, i // 16, dtype=torch.float8_e4m3fn) + self.w2_weight_scale_2 = param(e, dtype=torch.float32) + self.register_buffer("route_counter", make_route_counter(), persistent=False) + self._loaded_names: set[str] = set() + self._fork = torch.cuda.Event() + self._join = torch.cuda.Event() + self.weights: FlintMoEWeights | None = None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + inter = self.w2_weight.shape[2] * 2 + for name, w in weights: + m = _EXPERT_PARAM.fullmatch(name) + if m is None: + raise ValueError(f"Unexpected Flint expert weight: {name}") + e, proj, kind = int(m.group(1)), m.group(2), m.group(3) + param_name = ("w2_" if proj == "down_proj" else "w13_") + kind + param = getattr(self, param_name) + if proj == "down_proj": + param.data[e].copy_(w) + elif kind == "weight_scale_2": + param.data[e, int(proj == "up_proj")].copy_(w) + else: + rows = slice(inter, 2 * inter) if proj == "up_proj" else slice(inter) + param.data[e, rows].copy_(w) + loaded.add(param_name) + self._loaded_names.add(name) + if len(self._loaded_names) == 9 * self.num_experts: + self._repack() + return loaded + + def _repack(self) -> None: + gs = self.w13_weight_scale_2.data + if not torch.equal(gs[:, 0], gs[:, 1]): + raise ValueError("Flint NVFP4 gate/up global scales must match") + self.weights = FlintMoEWeights.from_modelopt( + self.w13_weight.data, + self.w13_weight_scale.data, + gs[:, 0], + self.w2_weight.data, + self.w2_weight_scale.data, + self.w2_weight_scale_2.data, + ) + for name in list(self._parameters): + delattr(self, name) + if self.max_num_tokens > DECODE_MAX_TOKENS: + moe_cute.warmup( + self.weights, self.top_k, DECODE_MAX_TOKENS + 1, self.max_num_tokens + ) + + def forward( + self, + hidden_states: torch.Tensor, + router_weight: torch.Tensor, + shared: torch.Tensor | nn.Module, + ) -> torch.Tensor: + """shared: shared-expert output, or the shared-expert module, which + then runs on the auxiliary stream concurrently with the routed path.""" + if self.weights is None: + raise RuntimeError("Flint NVFP4 expert weights were not fully loaded") + shared, ready = _start_shared(shared, hidden_states, self._fork, self._join) + return flint_moe( + hidden_states, + router_weight, + self.weights, + self.top_k, + self.route_counter, + shared, + ready, + ) + + +# Low-latency decode MoE block: each kernel prefetches weights before its PDL wait. + +# Largest batch for which the fused decode MoE block beats the regular path. +LL_MOE_MAX_TOKENS = 2 + + +@dataclass +class LLMoEWeights: + """Weights of one NVFP4 MoE layer for ``flint_moe_ll`` (views, no copies). + + Shared expert: ModelOpt-packed weights (int32 views) with flat + 128x4-swizzled scale bytes and fp32 global scales. + """ + + router: torch.Tensor # bf16 [E, H] + experts: FlintMoEWeights + shared_gu: torch.Tensor # [2 * SI, H / 8] int32 + shared_gu_scale: torch.Tensor + shared_gu_gs: torch.Tensor + shared_down: torch.Tensor # [H, SI / 8] int32 + shared_down_scale: torch.Tensor + shared_down_gs: torch.Tensor + + +def flint_moe_ll( + attn_out: torch.Tensor, + residual: torch.Tensor, + ln_weight: torch.Tensor, + eps: float, + w: LLMoEWeights, + top_k: int, + counter: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Decode MoE block for M <= LL_MOE_MAX_TOKENS tokens. + + Computes ``residual' = attn_out + residual`` and ``x = rms_norm(residual')`` + as ``norm.fused_add_rms_norm``, then the routed + shared expert MoE of + ``x`` (NVFP4 rule). Returns ``(moe_out, residual')``; inputs are not + modified. A Triton router kernel and two CuTe split-K expert kernels + (``moe_gemv``) chained with programmatic dependent launch. counter: buffer + from ``make_route_counter``. + """ + m, hidden = attn_out.shape + assert 0 < m <= LL_MOE_MAX_TOKENS and attn_out.is_contiguous() + assert residual.is_contiguous() + num_experts = w.router.shape[0] + si = w.shared_down.shape[1] * 8 + ew = w.experts + dev = attn_out.device + bf16 = torch.bfloat16 + res_out = torch.empty_like(residual) + xn = torch.empty_like(attn_out) + logits = torch.empty(m, num_experts, dtype=bf16, device=dev) + ids = torch.empty(m, top_k, dtype=torch.int32, device=dev) + tw = torch.empty(m, top_k, dtype=torch.float32, device=dev) + sact = torch.empty(m, si, dtype=bf16, device=dev) + act = torch.empty(m * top_k, ew.intermediate_size, dtype=bf16, device=dev) + out = torch.empty_like(attn_out) + eb = 2 + _router_kernel[(m * (num_experts // eb),)]( + attn_out, + residual, + ln_weight, + eps, + res_out, + xn, + w.router, + logits, + ids, + tw, + counter, + H=hidden, + E=num_experts, + TOPK=top_k, + EB=eb, + FUSE_NORM=True, + num_warps=2, + launch_pdl=True, + ) + moe_gemv.gate_up( + xn, + ids, + ew.w1, + ew.w1_sf, + ew.w1_alpha, + act, + w.shared_gu, + w.shared_gu_scale, + w.shared_gu_gs, + sact, + ) + moe_gemv.down_combine( + act, + ids, + tw, + ew.w2, + ew.w2_sf, + ew.w2_alpha, + sact, + w.shared_down.view(torch.uint8), + w.shared_down_scale, + w.shared_down_gs, + out, + ) + return out, res_out + + +def ll_weights_from_layers( + router_weight: torch.Tensor, + shared_gate_up: nn.Module, + shared_down: nn.Module, + experts: FlintMoEWeights, +) -> LLMoEWeights: + """``LLMoEWeights`` viewing the loaded layers' tensors: the shared-expert + linears are vLLM linear layers holding the FlashInfer-prepared NVFP4 + tensors (ModelOpt-packed weight, 128x4-swizzled scales, fp32 global + scale).""" + + def nvfp4(layer: nn.Module) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + weight = layer.weight + n, k2 = weight.shape + k = k2 * 2 + scale = layer.weight_scale + nbytes = (n + 127) // 128 * 128 * ((k // 16 + 3) // 4 * 4) + storage = scale.untyped_storage() + assert ( + weight.dtype == torch.uint8 + and weight.is_contiguous() + and not getattr(layer, "weights_padding_cols", 0) + and scale.storage_offset() == 0 + and storage.nbytes() == nbytes + and n % 128 == 0 + and k % 64 == 0 + ), "Unexpected NVFP4 shared-expert weight layout" + flat = torch.empty(0, dtype=torch.uint8, device=weight.device).set_( + storage, 0, (nbytes,), (1,) + ) + gs = layer.weight_global_scale.reshape(1).float() + return weight.view(torch.int32), flat, gs + + gu_w, gu_s, gu_gs = nvfp4(shared_gate_up) + dn_w, dn_s, dn_gs = nvfp4(shared_down) + return LLMoEWeights( + router=router_weight, + experts=experts, + shared_gu=gu_w, + shared_gu_scale=gu_s, + shared_gu_gs=gu_gs, + shared_down=dn_w, + shared_down_scale=dn_s, + shared_down_gs=dn_gs, + ) diff --git a/vllm/models/flint/ops/moe_cute.py b/vllm/models/flint/ops/moe_cute.py new file mode 100644 index 00000000000..c44753a3ef1 --- /dev/null +++ b/vllm/models/flint/ops/moe_cute.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""FlashInfer's SM100 NVFP4-weight / BF16-activation grouped MoE GEMMs. + +The FlashInfer kernel decodes E2M1 * E4M3 weight tiles to BF16 and runs BF16 +``tcgen05`` MMAs with fp32 accumulation and the per-expert fp32 global scale +applied to the accumulator. Its SwiGLU epilogue keeps gate/up in fp32 and +rounds ``silu(g) * u`` once (sigmoid via ``ex2.approx`` + ``rcp.approx``), the +rounding rule every Flint MoE path follows. +""" + +from functools import cache + +import torch + + +# ((mma_m, route_tile, mma_k), cluster_mn, raster_along_m) for GEMM1 / GEMM2, +# tuned on GB200 for E=256, H=2048, I=512, top-6. +def _tactic(num_tokens: int): + if num_tokens >= 4096: + return (((256, 128, 256), (2, 1), True), ((128, 128, 256), (2, 1), True)) + if num_tokens >= 2048: + return (((256, 128, 256), (2, 1), True), ((256, 128, 256), (2, 1), True)) + if num_tokens >= 512: + return (((128, 32, 256), (2, 1), True), ((128, 32, 256), (2, 1), True)) + return (((128, 8, 256), (2, 1), True), ((128, 8, 256), (2, 1), True)) + + +def grouped_moe( + x: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + w1: torch.Tensor, + w1_sf: torch.Tensor, + w1_alpha: torch.Tensor, + w2: torch.Tensor, + w2_sf: torch.Tensor, + w2_alpha: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Routed expert outputs in permuted order (unweighted, BF16). + + Returns (y [route_slots, H], expanded_idx_to_permuted_idx [M, TOPK]); the + caller does the weighted top-k reduction. + """ + from flashinfer.fused_moe.cute_dsl.blackwell import moe_w4a16 as lmod + from flashinfer.fused_moe.cute_dsl.moe_utils import ( + allocate_moe_sort_buffers, + get_max_num_permuted_tokens, + moe_permute, + moe_sort, + ) + from flashinfer.tllm_enums import ( + DEFAULT_SWIGLU_ALPHA, + DEFAULT_SWIGLU_BETA, + DEFAULT_SWIGLU_LIMIT, + ActivationType, + ) + + num_tokens, hidden = x.shape + top_k = topk_ids.shape[1] + num_experts = w1.shape[0] + inter = w2.shape[2] * 2 + gemm1_tactic, gemm2_tactic = _tactic(num_tokens) + route_tile = gemm1_tactic[0][1] + route_slots = get_max_num_permuted_tokens( + num_tokens, top_k, num_experts, route_tile + ) + sort_bufs = allocate_moe_sort_buffers( + num_tokens=num_tokens, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_experts, + tile_tokens_dim=route_tile, + device=x.device, + ) + ( + tile_to_expert, + tile_to_mn_limit, + expanded_to_permuted, + permuted_to_expanded, + _, + num_tiles, + ) = moe_sort( + token_selected_experts=topk_ids, + token_final_scales=topk_weights, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + num_local_experts=num_experts, + tile_tokens_dim=route_tile, + enable_pdl=True, + **sort_bufs, + ) + num_route_tiles = route_slots // route_tile + tile_to_expert = tile_to_expert[:num_route_tiles] + tile_to_mn_limit = tile_to_mn_limit[:num_route_tiles] + permuted_to_expanded = permuted_to_expanded[:route_slots] + hidden_ws = torch.empty(route_slots, hidden, dtype=x.dtype, device=x.device) + inter_ws = torch.empty(route_slots, inter, dtype=x.dtype, device=x.device) + moe_permute( + input=x, + permuted_output=hidden_ws, + tile_idx_to_mn_limit=tile_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_to_expanded, + num_non_exiting_tiles=num_tiles, + max_num_permuted_tokens=route_slots, + top_k=top_k, + tile_size=route_tile, + enable_pdl=True, + ) + common = dict( + tile_idx_to_expert_idx=tile_to_expert, + tile_idx_to_mn_limit=tile_to_mn_limit, + num_non_exiting_tiles=num_tiles, + num_local_experts=num_experts, + swiglu_alpha=DEFAULT_SWIGLU_ALPHA, + swiglu_beta=DEFAULT_SWIGLU_BETA, + swiglu_limit=DEFAULT_SWIGLU_LIMIT, + situ_beta=None, + situ_linear_beta=None, + use_fused_finalize=False, + permuted_idx_to_expanded_idx=None, + token_final_scales=None, + enable_pdl=True, + ) + lmod._run_grouped_gemm( + weight=w1, + weight_sf=w1_sf, + activations=hidden_ws, + alpha=w1_alpha, + output=inter_ws, + activation_type=ActivationType.Swiglu, + tactic=gemm1_tactic, + **common, + ) + # GEMM2 writes over the permuted activations, which are dead by now. + lmod._run_grouped_gemm( + weight=w2, + weight_sf=w2_sf, + activations=inter_ws, + alpha=w2_alpha, + output=hidden_ws, + activation_type=None, + tactic=gemm2_tactic, + **common, + ) + return hidden_ws, expanded_to_permuted + + +def _cdiv(a: int, b: int) -> int: + return -(-a // b) + + +def _compile_key(num_tokens: int, top_k: int, num_experts: int, gemm1_rows: int): + """The parts of the kernel cache keys that vary with num_tokens (tactics + and whether GEMM1 uses the CLC scheduler, as in FlashInfer's launcher).""" + from flashinfer.cute_dsl.utils import get_max_active_clusters + from flashinfer.fused_moe.cute_dsl.moe_utils import get_max_num_permuted_tokens + + gemm1, gemm2 = _tactic(num_tokens) + (mma_m, route_tile, _), cluster, _ = gemm1 + slots = get_max_num_permuted_tokens(num_tokens, top_k, num_experts, route_tile) + cta_m = mma_m // (2 if mma_m == 256 else 1) + clusters = _cdiv(_cdiv(gemm1_rows, cta_m), cluster[0]) * _cdiv( + _cdiv(slots, route_tile), cluster[1] + ) + clc = clusters > get_max_active_clusters(cluster[0] * cluster[1]) + return gemm1, gemm2, clc + + +@cache +def _variants(top_k, num_experts, gemm1_rows, min_tokens, max_tokens) -> dict: + """Smallest token count of each kernel cache key in [min, max].""" + reps: dict = {} + for n in range(min_tokens, max_tokens + 1): + reps.setdefault(_compile_key(n, top_k, num_experts, gemm1_rows), n) + return reps + + +_WARMED: set[tuple] = set() + + +def warmup(weights, top_k: int, min_tokens: int, max_tokens: int) -> None: + """Compile every grouped-GEMM kernel variant used for min..max tokens. + + CuTe kernels compile on first use, which must not happen inside CUDA graph + capture; the compile caches are process-wide, so each variant runs once. + """ + num_experts, gemm1_rows = weights.w1.shape[:2] + hidden = weights.w2.shape[1] + reps = _variants(top_k, num_experts, gemm1_rows, min_tokens, max_tokens) + for key, n in reps.items(): + if key in _WARMED: + continue + x = torch.zeros(n, hidden, dtype=torch.bfloat16, device=weights.w1.device) + ids = torch.randint(0, num_experts, (n, top_k), device=x.device) + tw = torch.full((n, top_k), 1.0 / top_k, device=x.device) + grouped_moe( + x, + ids.int(), + tw, + weights.w1, + weights.w1_sf, + weights.w1_alpha, + weights.w2, + weights.w2_sf, + weights.w2_alpha, + ) + _WARMED.add(key) + torch.accelerator.synchronize() diff --git a/vllm/models/flint/ops/moe_gemv.py b/vllm/models/flint/ops/moe_gemv.py new file mode 100644 index 00000000000..467801cee33 --- /dev/null +++ b/vllm/models/flint/ops/moe_gemv.py @@ -0,0 +1,573 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Expert-indexed small-M NVFP4 GEMV kernels (CuTe DSL) for the low-latency +decode MoE block, built on the dense GEMV design of ``nvfp4_gemv``: each CTA +owns 16 output rows (the m16n8k16 M dimension, swap-AB), its warps split K and +reduce through shared memory, weights are dequantized exactly to BF16 +(``cvt.rn.bf16x2.e2m1x2`` + ``mul.rn.bf16x2``) and the fp32 global scale is +applied to the fp32 accumulator. + +``gate_up``: one (token, expert) pair per CTA column of the grid; writes +``bf16(silu(g) * u)`` from fp32 gate/up with FlashInfer's approximate sigmoid +(the shared expert rounds its gate/up to BF16 first and uses the exact SiLU, +as its dense GEMM does). Expert weights use FlashInfer's W4A16 +layout ([up; gate] in 64-row groups, 6D MMA scales). + +``down_combine``: per token and 16 output rows, the down GEMVs of the token's +top-k experts and of the shared expert, all weights prefetched before +``griddepcontrol.wait``; ``out = bf16(bf16(sum_j tw_j * y_j) + y_shared)`` +with ``y = bf16(acc * global_scale)``. +""" + +from typing import Any + +import torch + +_CHUNK = 64 +_WARPS = 4 +_SPLIT = 2 + +_classes: Any = None +_compiled: dict[tuple, Any] = {} + + +def _build(): + import cutlass + import cutlass.cute as cute + from cuda.bindings.driver import CUstream + from cutlass import Float32, Int32, Int64, Uint32 + from cutlass.cutlass_dsl import dsl_user_op + + from .cute_primitives import ( + alloc_reduce, + cta_reduce, + decode, + ld_x16, + load_rows, + mma, + ptx, + round_bf16, + scale_pair, + silu_mul, + st_out, + ) + + @dsl_user_op + def silu_mul_approx(g: Float32, u: Float32, *, loc=None, ip=None): + # g * rcp.approx(1 + ex2.approx(-g * log2(e))) * u, the instructions of + # FlashInfer's W4A16 MoE epilogue (routed experts). + return ptx( + "{\n.reg .f32 e, d;\n" + "mul.f32 e, $1, 0fBFB8AA3B;\nex2.approx.ftz.f32 e, e;\n" + "add.f32 d, e, 0f3F800000;\nrcp.approx.ftz.f32 d, d;\n" + "mul.f32 d, $1, d;\nmul.f32 $0, d, $2;\n}", + "=f,f,f", + [g, u], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + @cute.jit + def gemv_pass( + acc, + a_off: cutlass.Constexpr, + qs, + ss, + q_off: cutlass.Constexpr, + s_off: cutlass.Constexpr, + x_base, + CPW: cutlass.Constexpr, + valid, + ): + """acc[a_off:a_off+4] += W_rows (16 x 64*CPW) @ x (token column 0).""" + for c in cutlass.range_constexpr(CPW): + xv = ld_x16(x_base + c * 128, valid) + sg = scale_pair(ss[s_off + c * 2]) + sg8 = scale_pair(ss[s_off + c * 2 + 1]) + q0 = q_off + c * 4 + for j in cutlass.range_constexpr(4): + a = decode(qs[q0 + j // 2], qs[q0 + 2 + j // 2], j, sg, sg8) + r = mma( + [acc[a_off], acc[a_off + 1], acc[a_off + 2], acc[a_off + 3]], + a, + xv[2 * j], + xv[2 * j + 1], + ) + for i in cutlass.range_constexpr(4): + acc[a_off + i] = r[i] + + class GateUp: + def __init__(self, h, inter, si, topk): + self.h, self.inter, self.si, self.topk = h, inter, si, topk + assert si % inter == 0 + self.kc = h // _CHUNK + assert self.kc % (_WARPS * _SPLIT) == 0 + self.cpw = self.kc // (_WARPS * _SPLIT) + + @cute.jit + def __call__( + self, + x: cute.Tensor, + ids: cute.Tensor, + w: cute.Tensor, + s: cute.Tensor, + gs: cute.Tensor, + act: cute.Tensor, + wsh: cute.Tensor, + ssh: cute.Tensor, + gssh: cute.Tensor, + sact: cute.Tensor, + P: Int32, + NSH: Int32, + stream: CUstream, + ): + self.kernel(x, ids, w, s, gs, act, wsh, ssh, gssh, sact, P).launch( + grid=[self.inter // 16 * _SPLIT, P + NSH, 1], + block=[32 * _WARPS, 1, 1], + cluster=[_SPLIT, 1, 1], + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + ids: cute.Tensor, + w: cute.Tensor, + s: cute.Tensor, + gs: cute.Tensor, + act: cute.Tensor, + wsh: cute.Tensor, + ssh: cute.Tensor, + gssh: cute.Tensor, + sact: cute.Tensor, + P: Int32, + ): + H: cutlass.Constexpr = self.h + INTER: cutlass.Constexpr = self.inter + KC: cutlass.Constexpr = self.kc + CPW: cutlass.Constexpr = self.cpw + WARPS: cutlass.Constexpr = _WARPS + tidx, _, _ = cute.arch.thread_idx() + SPLIT: cutlass.Constexpr = _SPLIT + bx, p, _ = cute.arch.block_idx() + nb = bx // SPLIT + rank = bx % SPLIT + warp = tidx // 32 + lane = tidx % 32 + g = lane // 4 + t = lane % 4 + c_begin = (rank * WARPS + warp) * CPW + SI: cutlass.Constexpr = self.si + # Programs p >= P are the shared expert: token (p - P) // halves, + # columns ((p - P) % halves) * INTER + 16 * nb. + HALVES: cutlass.Constexpr = self.si // self.inter + row_g = Int32(0) + row_u = Int32(0) + w_base = Int64(0) + s_base = Int64(0) + gsv = Float32(0.0) + tok = Int32(0) + o_off = Int64(0) + qs = cute.make_rmem_tensor((2 * CPW * 4,), Uint32) + ss = cute.make_rmem_tensor((2 * CPW * 2,), Uint32) + if p < P: + cute.arch.griddepcontrol_wait() + e = Int64(ids[p]) + # FlashInfer GEMM1 rows: [up 64; gate 64] per 128-row group. + row_u = (nb * 16) // 64 * 128 + (nb * 16) % 64 + g + row_g = row_u + 64 + w_base = w.iterator.toint() + e * (2 * INTER * KC * 32) + s_base = s.iterator.toint() + e * ((2 * INTER // 128) * KC * 512) + gsv = gs[e] + tok = p // self.topk + o_off = Int64(p) * INTER + nb * 16 + load_rows( + w_base, s_base, row_g, KC, CPW, c_begin, t, qs, ss, 0, 0, True + ) + load_rows( + w_base, + s_base, + row_u, + KC, + CPW, + c_begin, + t, + qs, + ss, + CPW * 4, + CPW * 2, + True, + ) + else: + # Shared expert: static weights, loaded before the wait. + v = p - P + tok = v // HALVES + c0 = (v % HALVES) * INTER + nb * 16 + row_g = c0 + g + row_u = SI + c0 + g + w_base = wsh.iterator.toint() + s_base = ssh.iterator.toint() + gsv = gssh[0] + o_off = Int64(tok) * SI + c0 + load_rows( + w_base, s_base, row_g, KC, CPW, c_begin, t, qs, ss, 0, 0, True + ) + load_rows( + w_base, + s_base, + row_u, + KC, + CPW, + c_begin, + t, + qs, + ss, + CPW * 4, + CPW * 2, + True, + ) + cute.arch.griddepcontrol_wait() + cute.arch.cluster_arrive_relaxed() + x_base = x.iterator.toint() + (Int64(tok) * H + c_begin * 64 + t * 16) * 2 + valid = Int32(g == 0) + acc = cute.make_rmem_tensor((8,), Float32) + acc.fill(0.0) + gemv_pass(acc, 0, qs, ss, 0, 0, x_base, CPW, valid) + gemv_pass(acc, 4, qs, ss, CPW * 4, CPW * 2, x_base, CPW, valid) + cute.arch.griddepcontrol_launch_dependents() + smem = cutlass.memory.SmemAllocator() + red, part = alloc_reduce(smem, 8, WARPS, SPLIT) + cta_reduce(acc, 8, WARPS, SPLIT, warp, lane, rank, red, part) + if warp == 0 and rank == 0 and t == 0: + a_base = act.iterator.toint() + if p >= P: + a_base = sact.iterator.toint() + # Routed experts keep gate/up in fp32 (TRTLLM MoE rule); the + # shared expert is a dense MLP whose gate/up GEMM output is BF16. + is_shared = p >= P + for i in cutlass.range_constexpr(0, 4, 2): + gv = acc[i] * gsv + uv = acc[4 + i] * gsv + gv = cutlass.select_(is_shared, round_bf16(gv), gv) + uv = cutlass.select_(is_shared, round_bf16(uv), uv) + v = cutlass.select_( + is_shared, silu_mul(gv, uv), silu_mul_approx(gv, uv) + ) + off = o_off + g + 8 * (i // 2) + st_out(a_base + off * 2, v, Int32(1)) + + class DownCombine: + def __init__(self, h, inter, si, topk): + self.h, self.inter, self.si, self.topk = h, inter, si, topk + self.kc = inter // _CHUNK + self.kcs = si // _CHUNK + parts = _WARPS * _SPLIT + assert self.kc % parts == 0 and self.kcs % parts == 0 + self.cpw = self.kc // parts + self.cpws = self.kcs // parts + + @cute.jit + def __call__( + self, + act: cute.Tensor, + ids: cute.Tensor, + tw: cute.Tensor, + w: cute.Tensor, + s: cute.Tensor, + gs: cute.Tensor, + sact: cute.Tensor, + wsd: cute.Tensor, + ssd: cute.Tensor, + gssd: cute.Tensor, + out: cute.Tensor, + M: Int32, + stream: CUstream, + ): + self.kernel(act, ids, tw, w, s, gs, sact, wsd, ssd, gssd, out).launch( + grid=[self.h // 16 * _SPLIT, M, 1], + block=[32 * _WARPS, 1, 1], + cluster=[_SPLIT, 1, 1], + stream=stream, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + act: cute.Tensor, + ids: cute.Tensor, + tw: cute.Tensor, + w: cute.Tensor, + s: cute.Tensor, + gs: cute.Tensor, + sact: cute.Tensor, + wsd: cute.Tensor, + ssd: cute.Tensor, + gssd: cute.Tensor, + out: cute.Tensor, + ): + H: cutlass.Constexpr = self.h + INTER: cutlass.Constexpr = self.inter + SI: cutlass.Constexpr = self.si + TOPK: cutlass.Constexpr = self.topk + KC: cutlass.Constexpr = self.kc + KCS: cutlass.Constexpr = self.kcs + CPW: cutlass.Constexpr = self.cpw + CPWS: cutlass.Constexpr = self.cpws + WARPS: cutlass.Constexpr = _WARPS + tidx, _, _ = cute.arch.thread_idx() + SPLIT: cutlass.Constexpr = _SPLIT + bx, tok, _ = cute.arch.block_idx() + nb = bx // SPLIT + rank = bx % SPLIT + warp = tidx // 32 + lane = tidx % 32 + g = lane // 4 + t = lane % 4 + part_id = rank * WARPS + warp + row = nb * 16 + g + qs = cute.make_rmem_tensor(((TOPK * CPW + CPWS) * 4,), Uint32) + ss = cute.make_rmem_tensor(((TOPK * CPW + CPWS) * 2,), Uint32) + # Expert ids / weights come from the router stage, which completed + # before the previous kernel passed its dependency wait. + gsv = cute.make_rmem_tensor((TOPK,), Float32) + twv = cute.make_rmem_tensor((TOPK,), Float32) + for j in cutlass.range_constexpr(TOPK): + e = Int64(ids[tok * TOPK + j]) + load_rows( + w.iterator.toint() + e * (H * KC * 32), + s.iterator.toint() + e * ((H // 128) * KC * 512), + row, + KC, + CPW, + part_id * CPW, + t, + qs, + ss, + j * CPW * 4, + j * CPW * 2, + True, + ) + gsv[j] = gs[e] + twv[j] = tw[tok * TOPK + j] + load_rows( + wsd.iterator.toint(), + ssd.iterator.toint(), + row, + KCS, + CPWS, + part_id * CPWS, + t, + qs, + ss, + TOPK * CPW * 4, + TOPK * CPW * 2, + True, + ) + gss = gssd[0] + cute.arch.cluster_arrive_relaxed() + cute.arch.griddepcontrol_wait() + valid = Int32(g == 0) + NA: cutlass.Constexpr = (TOPK + 1) * 4 + acc = cute.make_rmem_tensor((NA,), Float32) + acc.fill(0.0) + for j in cutlass.range_constexpr(TOPK): + x_base = ( + act.iterator.toint() + + (Int64(tok * TOPK + j) * INTER + part_id * CPW * 64 + t * 16) * 2 + ) + gemv_pass( + acc, j * 4, qs, ss, j * CPW * 4, j * CPW * 2, x_base, CPW, valid + ) + x_base = ( + sact.iterator.toint() + + (Int64(tok) * SI + part_id * CPWS * 64 + t * 16) * 2 + ) + gemv_pass( + acc, + TOPK * 4, + qs, + ss, + TOPK * CPW * 4, + TOPK * CPW * 2, + x_base, + CPWS, + valid, + ) + cute.arch.griddepcontrol_launch_dependents() + smem = cutlass.memory.SmemAllocator() + red, part = alloc_reduce(smem, NA, WARPS, SPLIT) + cta_reduce(acc, NA, WARPS, SPLIT, warp, lane, rank, red, part) + if warp == 0 and rank == 0 and t == 0: + o_base = out.iterator.toint() + for i in cutlass.range_constexpr(0, 4, 2): + routed = Float32(0.0) + for j in cutlass.range_constexpr(TOPK): + routed = routed + twv[j] * round_bf16(acc[j * 4 + i] * gsv[j]) + shared = round_bf16(acc[TOPK * 4 + i] * gss) + v = round_bf16(routed) + shared + r = nb * 16 + g + 8 * (i // 2) + st_out(o_base + (Int64(tok) * H + r) * 2, v, Int32(1)) + + return GateUp, DownCombine + + +def _fakes(*specs) -> list: + """Compile-time tensors: (dtype, divisibility) is 1-D of symbolic size, + (dtype, divisibility, shape) has that shape.""" + import cutlass + import cutlass.cute as cute + from quack.compile_utils import make_fake_tensor + + return [ + make_fake_tensor( + getattr(cutlass, spec[0]), + spec[2] if len(spec) > 2 else (cute.sym_int(),), + divisibility=spec[1], + ) + for spec in specs + ] + + +def _flat_bytes(t: torch.Tensor) -> torch.Tensor: + """The contiguous storage behind a (possibly permuted) scale view.""" + t = t.view(torch.uint8) + return torch.as_strided(t, (t.numel(),), (1,), t.storage_offset()) + + +def _stream(): + from cuda.bindings.driver import CUstream + + return CUstream(torch.cuda.current_stream().cuda_stream) + + +def _compile(key: tuple, fakes: list, *scalars): + """Compile kernel ``key = (kind, h, inter, si, topk)`` into ``_compiled``.""" + global _classes + import cutlass.cute as cute + + if _classes is None: + _classes = _build() + kern = _classes[key[0] == "down"](*key[1:]) + fn = _compiled[key] = cute.compile( + kern, *fakes, *scalars, _stream(), options="--enable-tvm-ffi" + ) + return fn + + +def gate_up( + x: torch.Tensor, + ids: torch.Tensor, + w1: torch.Tensor, + w1_sf: torch.Tensor, + alpha: torch.Tensor, + act: torch.Tensor, + wsh: torch.Tensor, + ssh: torch.Tensor, + gssh: torch.Tensor, + sact: torch.Tensor, +) -> None: + """act[p] = silu_and_mul((x[p // topk] @ W1[ids[p]]^T) * alpha) on fp32 + gate/up for the FlashInfer-layout NVFP4 experts (w1 uint8 [E, 2I, H/2]), + and in the same launch sact[t] = silu_and_mul(bf16((x[t] @ Wsh^T) * gssh)) + for the shared expert (ModelOpt [gate; up] uint8 [2 SI, H/2], flat + 128x4-swizzled scales).""" + m, h = x.shape + p, topk = ids.numel(), ids.shape[1] + inter = w1.shape[1] // 2 + si = wsh.shape[0] // 2 + key = ("gu", h, inter, si, topk) + fn = _compiled.get(key) + if fn is None: + fn = _compile( + key, + _fakes( + ("BFloat16", 8), + ("Int32", 1), + ("Uint8", 16), + ("Uint8", 4), + ("Float32", 1), + ("BFloat16", 1), + ("Uint8", 16), + ("Uint8", 4), + ("Float32", 1, (1,)), + ("BFloat16", 1), + ), + 1, + 1, + ) + fn( + x.view(-1), + ids.view(-1), + w1.view(-1), + _flat_bytes(w1_sf), + alpha, + act.view(-1), + wsh.view(torch.uint8).view(-1), + ssh.view(-1), + gssh.reshape(1), + sact.view(-1), + p, + m * (si // inter), + _stream(), + ) + + +def down_combine( + act: torch.Tensor, + ids: torch.Tensor, + tw: torch.Tensor, + w2: torch.Tensor, + w2_sf: torch.Tensor, + alpha2: torch.Tensor, + sact: torch.Tensor, + wsd: torch.Tensor, + ssd: torch.Tensor, + gssd: torch.Tensor, + out: torch.Tensor, +) -> None: + """out[t] = bf16(bf16(sum_j tw * bf16(act_j @ W2[e_j]^T * alpha2)) + + bf16(sact[t] @ Wsd^T * gssd)).""" + m, topk = ids.shape + h = w2.shape[1] + inter = w2.shape[2] * 2 + si = wsd.shape[1] * 2 + key = ("down", h, inter, si, topk) + fn = _compiled.get(key) + if fn is None: + fn = _compile( + key, + _fakes( + ("BFloat16", 8), + ("Int32", 1), + ("Float32", 1), + ("Uint8", 16), + ("Uint8", 4), + ("Float32", 1), + ("BFloat16", 8), + ("Uint8", 16), + ("Uint8", 4), + ("Float32", 1, (1,)), + ("BFloat16", 1), + ), + 1, + ) + fn( + act.view(-1), + ids.view(-1), + tw.view(-1), + w2.view(-1), + _flat_bytes(w2_sf), + alpha2, + sact.view(-1), + wsd.view(-1), + ssd.view(-1), + gssd.reshape(1), + out.view(-1), + m, + _stream(), + ) diff --git a/vllm/models/flint/ops/norm.py b/vllm/models/flint/ops/norm.py new file mode 100644 index 00000000000..507c9b334c3 --- /dev/null +++ b/vllm/models/flint/ops/norm.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flint's RMSNorm, fused residual-add RMSNorm and SiLU-and-mul (Triton). + +Numerics of the reference implementation's CUDA norm and activation kernels, +everything in fp32 with one output rounding: + +- ``rms_norm``: ``bf16(x * rsqrt(mean(x^2) + eps) * w)``. +- ``fused_add_rms_norm``: ``s = x + residual`` in fp32, + ``residual_out = bf16(s)``, ``out = bf16(s * rsqrt(mean(s^2) + eps) * w)`` + from the unrounded ``s``. +- ``silu_and_mul``: ``bf16(g / (1 + exp(-g)) * u)`` (precise ``expf`` and + division). + +``rsqrt`` is ``rsqrt.approx.ftz.f32`` (``rsqrtf`` under fast math). The fused +Flint kernels (``qkv_fused``, the decode MoE block's norm and shared expert, and +the GEMV SiLU epilogue) follow the same rules; the routed experts use the +approximate ``ex2`` / ``rcp`` SiLU of the TRTLLM MoE (see ``moe``). +""" + +import torch +from torch import nn + +from vllm.triton_utils import tl, tldevice, triton + + +@triton.jit +def _rsqrt(x): + return tl.inline_asm_elementwise( + "rsqrt.approx.ftz.f32 $0, $1;", "=f,f", [x], tl.float32, True, 1 + ) + + +@triton.jit +def _rms_norm_kernel( + x_ptr, + res_ptr, + w_ptr, + out_ptr, + res_out_ptr, + x_stride, + eps, + H: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, +): + row = tl.program_id(0).to(tl.int64) + offs = tl.arange(0, H) + w = tl.load(w_ptr + offs).to(tl.float32) + s = tl.load(x_ptr + row * x_stride + offs).to(tl.float32) + if HAS_RESIDUAL: + s += tl.load(res_ptr + row * H + offs).to(tl.float32) + tl.store(res_out_ptr + row * H + offs, s.to(tl.bfloat16)) + rs = _rsqrt(tl.sum(s * s, axis=0) / H + eps) + tl.store(out_ptr + row * H + offs, (s * rs * w).to(tl.bfloat16)) + + +def _launch_norm(x, residual, weight, eps, out): + m, h = x.shape + assert x.dtype == torch.bfloat16 and x.stride(-1) == 1 + assert weight.shape == (h,) and h & (h - 1) == 0 + if m > 0: + _rms_norm_kernel[(m,)]( + x, + x if residual is None else residual, + weight, + out, + residual, + x.stride(0), + eps, + H=h, + HAS_RESIDUAL=residual is not None, + num_warps=4, + ) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + _launch_norm(x, None, weight, eps, out) + return out + + +def fused_add_rms_norm( + x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float +) -> tuple[torch.Tensor, torch.Tensor]: + """In place: ``residual = bf16(x + residual)``, ``x = rms_norm(x + + residual)``; returns ``(x, residual)``.""" + assert residual.shape == x.shape and residual.is_contiguous() + assert x.is_contiguous() + _launch_norm(x, residual, weight, eps, x) + return x, residual + + +@triton.jit +def _silu_and_mul_kernel(x_ptr, out_ptr, x_stride, D: tl.constexpr, BD: tl.constexpr): + row = tl.program_id(0).to(tl.int64) + offs = tl.program_id(1) * BD + tl.arange(0, BD) + g = tl.load(x_ptr + row * x_stride + offs).to(tl.float32) + u = tl.load(x_ptr + row * x_stride + D + offs).to(tl.float32) + silu = tl.math.div_rn(g, 1.0 + tldevice.exp(-g)) + tl.store(out_ptr + row * D + offs, (silu * u).to(tl.bfloat16)) + + +def silu_and_mul(x: torch.Tensor) -> torch.Tensor: + """``silu(x[:, :D]) * x[:, D:]`` for ``x`` [M, 2D].""" + m, d2 = x.shape + d = d2 // 2 + assert x.dtype == torch.bfloat16 and x.stride(-1) == 1 + bd = min(d, 1024) + assert d % bd == 0 + out = torch.empty(m, d, dtype=x.dtype, device=x.device) + if m > 0: + _silu_and_mul_kernel[(m, d // bd)](x, out, x.stride(0), D=d, BD=bd, num_warps=4) + return out + + +class RMSNorm(nn.Module): + """Holds a norm ``weight``; ``rms_norm(x)``, or with a residual + ``fused_add_rms_norm(x, residual)``.""" + + def __init__(self, hidden_size: int, eps: float) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward( + self, x: torch.Tensor, residual: torch.Tensor | None = None + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if residual is None: + return rms_norm(x, self.weight, self.variance_epsilon) + return fused_add_rms_norm(x, residual, self.weight, self.variance_epsilon) diff --git a/vllm/models/flint/ops/nvfp4_gemv.py b/vllm/models/flint/ops/nvfp4_gemv.py new file mode 100644 index 00000000000..54d04ef107f --- /dev/null +++ b/vllm/models/flint/ops/nvfp4_gemv.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small-M BF16 x NVFP4 (W4A16) and BF16 x BF16 GEMM for SM100 in CuTe DSL. + +``out[M, N] = (x[M, K] @ dequant(W)[N, K]^T) * global_scale`` with the weight in +the ModelOpt layout used by the FlashInfer CuTe-DSL W4A16 kernel: ``weight`` +uint8 ``[N, K / 2]`` (two E2M1 codes per byte, low nibble first) and +``weight_scale`` the 128x4-swizzled E4M3 block scales (one per 16 values), so +both kernels share the same tensors. + +Numerics: every weight is dequantized exactly to BF16 (``e2m1 * e4m3`` via +``cvt.rn.bf16x2.e2m1x2`` and ``mul.rn.bf16x2``), multiplied with the BF16 +activations on tensor cores (``mma.sync`` m16n8k16, FP32 accumulate), and the +FP32 global scale is applied to the FP32 accumulator before the single output +rounding to BF16. + +Each CTA owns 16 output features (the MMA M dimension, swap-AB) and ``8 * nt`` +tokens (``nt`` = 1, 2 or 4, i.e. up to 32); its warps (and, with ``split > 1``, +the CTAs of a cluster) split K and reduce through shared memory / DSMEM. All +weight loads are issued before ``griddepcontrol.wait`` so they overlap the +previous kernel under PDL. +""" + +from dataclasses import dataclass +from typing import Any + +import torch + +_CHUNK = 64 # K values per warp step (four m16n8k16 k-tiles) +_NACC = 8 # accumulator sets per lane, rotated over k-tiles + + +@dataclass(frozen=True) +class GemvConfig: + warps: int + split: int + + +_C = GemvConfig +_Table = dict[tuple[int, int], dict[int, GemvConfig]] +# Tuned on GB200 (CUDA-graph replay, weights rotated past L2): config per +# number of 8-token MMA tiles per CTA (``nt``). A shape / ``nt`` is listed only +# where the GEMV beats the alternative (FlashInfer CuTe-DSL for NVFP4, cuBLAS +# for BF16) at ``M <= 8 * nt``; the model uses the GEMV only for listed ones. +# Gate/up projections always take the SiLU-and-mul tables. +_TUNED: _Table = { + (5120, 2048): {1: _C(8, 1), 2: _C(8, 2)}, + (2048, 4096): {1: _C(4, 4), 2: _C(8, 2), 4: _C(4, 4)}, + (2048, 6144): {1: _C(8, 4), 2: _C(8, 2), 4: _C(4, 8)}, + (2048, 1024): {1: _C(4, 1), 2: _C(4, 1), 4: _C(8, 1)}, +} +_TUNED_SILU_MUL: _Table = { + (12288, 2048): {1: _C(4, 1), 2: _C(2, 2)}, + (2048, 2048): {1: _C(2, 4), 2: _C(4, 2), 4: _C(8, 2)}, +} +# The BF16 qkv and dense down GEMVs are no faster than cuBLAS. +_TUNED_BF16: _Table = { + (2048, 4096): {1: _C(4, 8), 2: _C(8, 4)}, + (2048, 1024): {1: _C(4, 4), 2: _C(4, 4)}, +} +_TUNED_BF16_SILU_MUL: _Table = { + (12288, 2048): {1: _C(2, 1)}, + (2048, 2048): {1: _C(8, 4), 2: _C(4, 4), 4: _C(8, 2)}, +} + + +def _table(silu_mul: bool, bf16: bool) -> _Table: + if bf16: + return _TUNED_BF16_SILU_MUL if silu_mul else _TUNED_BF16 + return _TUNED_SILU_MUL if silu_mul else _TUNED + + +def _configs(n: int, k: int, silu_mul: bool, bf16: bool) -> dict[int, GemvConfig]: + return _table(silu_mul, bf16)[(n, k)] + + +def tuned(n: int, k: int, silu_mul: bool = False, bf16: bool = False) -> bool: + """Whether the GEMV is the fastest choice for this shape at some M.""" + return (n, k) in _table(silu_mul, bf16) + + +def pick_nt( + n: int, k: int, m: int, silu_mul: bool = False, bf16: bool = False +) -> int | None: + """Smallest tuned token tile count covering ``m``, or None if the GEMV + is not the fastest choice for this shape and M.""" + nts = [nt for nt in _table(silu_mul, bf16).get((n, k), ()) if 8 * nt >= m] + return min(nts) if nts else None + + +def _build_kernel_class(): + import cutlass + import cutlass.cute as cute + from cuda.bindings.driver import CUstream + from cutlass import Float32, Int32, Int64, Uint32, const_expr + + from .cute_primitives import ( + alloc_reduce, + cta_reduce, + decode, + ld_w16, + ld_x16, + load_rows, + mma, + round_bf16, + scale_pair, + silu_mul, + st_out, + ) + + class Nvfp4Gemv: + def __init__(self, k, nt, warps, split, pdl, silu_mul, bf16): + self.bf16 = bf16 + self.kc = k // _CHUNK + assert self.kc % (warps * split) == 0 + self.cpw = self.kc // (warps * split) + self.nt = nt + self.warps = warps + self.split = split + self.pdl = pdl + # silu_mul: each CTA also owns the matching "up" row block N/2 + # rows below and writes silu(gate) * up. + self.rb = 2 if silu_mul else 1 + self.nacc = _NACC // self.rb + + @cute.jit + def __call__( + self, + x: cute.Tensor, + w: cute.Tensor, + s: cute.Tensor, + gs: cute.Tensor, + out: cute.Tensor, + M: Int32, + N: Int32, + stride_x: Int32, + stride_o: Int32, + stream: CUstream, + ): + tokens = 8 * self.nt + n_out = N // self.rb + self.kernel(x, w, s, gs, out, M, n_out, stride_x, stride_o).launch( + grid=[(n_out // 16) * self.split, (M + tokens - 1) // tokens, 1], + block=[32 * self.warps, 1, 1], + cluster=[self.split, 1, 1], + stream=stream, + use_pdl=self.pdl, + ) + + @cute.kernel + def kernel( + self, + x: cute.Tensor, + w: cute.Tensor, + s: cute.Tensor, + gs: cute.Tensor, + out: cute.Tensor, + M: Int32, + n_out: Int32, + stride_x: Int32, + stride_o: Int32, + ): + NT: cutlass.Constexpr = self.nt + CPW: cutlass.Constexpr = self.cpw + KC: cutlass.Constexpr = self.kc + WARPS: cutlass.Constexpr = self.warps + SPLIT: cutlass.Constexpr = self.split + NACC: cutlass.Constexpr = self.nacc + RB: cutlass.Constexpr = self.rb + NA: cutlass.Constexpr = RB * NT * 4 # accumulators per lane + BF16W: cutlass.Constexpr = self.bf16 + KB: cutlass.Constexpr = self.kc * (128 if self.bf16 else 32) + QW: cutlass.Constexpr = 16 if self.bf16 else 4 # words per chunk + tidx, _, _ = cute.arch.thread_idx() + bx, mt, _ = cute.arch.block_idx() + nb = bx // SPLIT + rank = bx % SPLIT + warp = tidx // 32 + lane = tidx % 32 + g = lane // 4 + t = lane % 4 + c_begin = (rank * WARPS + warp) * CPW + + # Lane (g, t) of k-tile j covers k = 64c + 16t + 4j + {0..3} of + # rows g and g + 8, i.e. bytes 8t + 2j + {0, 1} of the 32-byte + # chunk row, and scale group 4c + t of both rows. + qs = cute.make_rmem_tensor((RB * CPW * QW,), Uint32) + ss = cute.make_rmem_tensor((RB * CPW * 2,), Uint32) + for rb in cutlass.range_constexpr(RB): + row_g = nb * 16 + g + rb * n_out + if const_expr(BF16W): + # 16 BF16 weights of rows g and g + 8 per chunk. + wbase = ( + w.iterator.toint() + Int64(row_g) * KB + c_begin * 128 + t * 32 + ) + for c in cutlass.range_constexpr(CPW): + q0 = (rb * CPW + c) * QW + for h in cutlass.range_constexpr(4): + off = (h // 2) * 8 * KB + c * 128 + (h % 2) * 16 + wv = ld_w16(wbase + off) + for i in cutlass.range_constexpr(4): + qs[q0 + h * 4 + i] = wv[i] + else: + load_rows( + w.iterator.toint(), + s.iterator.toint(), + row_g, + KC, + CPW, + c_begin, + t, + qs, + ss, + rb * CPW * QW, + rb * CPW * 2, + ) + gsv = Float32(1.0) + if const_expr(not BF16W): + gsv = gs[0] + if const_expr(SPLIT > 1): + # Peers must be running before DSMEM stores; waited on below. + cute.arch.cluster_arrive_relaxed() + + if const_expr(self.pdl): + cute.arch.griddepcontrol_wait() + xbase = x.iterator.toint() + (Int64(c_begin) * _CHUNK + t * 16) * 2 + xall = [] + for c in cutlass.range_constexpr(CPW): + xs = [] + for nt in cutlass.range_constexpr(NT): + tok = mt * (8 * NT) + nt * 8 + g + addr = xbase + Int64(tok) * stride_x * 2 + c * 128 + xs.append(ld_x16(addr, Int32(tok < M))) + xall.append(xs) + + acc = cute.make_rmem_tensor((NACC * NA,), Float32) + acc.fill(0.0) + for rb in cutlass.range_constexpr(RB): + for c in cutlass.range_constexpr(CPW): + q0 = (rb * CPW + c) * QW + if const_expr(not BF16W): + sg = scale_pair(ss[(rb * CPW + c) * 2]) + sg8 = scale_pair(ss[(rb * CPW + c) * 2 + 1]) + for j in cutlass.range_constexpr(4): + if const_expr(BF16W): + # rows g / g + 8, k pairs (4j, 4j+1) and (4j+2, 4j+3) + a = [ + qs[q0 + 2 * j], + qs[q0 + 8 + 2 * j], + qs[q0 + 2 * j + 1], + qs[q0 + 8 + 2 * j + 1], + ] + else: + a = decode(qs[q0 + j // 2], qs[q0 + 2 + j // 2], j, sg, sg8) + for nt in cutlass.range_constexpr(NT): + xv = xall[c][nt] + ab = ((c * 4 + j) % NACC) * NA + (rb * NT + nt) * 4 + r = mma( + [acc[ab], acc[ab + 1], acc[ab + 2], acc[ab + 3]], + a, + xv[2 * j], + xv[2 * j + 1], + ) + for i in cutlass.range_constexpr(4): + acc[ab + i] = r[i] + for h in cutlass.range_constexpr(1, NACC): + for i in cutlass.range_constexpr(NA): + acc[i] = acc[i] + acc[h * NA + i] + if const_expr(self.pdl): + cute.arch.griddepcontrol_launch_dependents() + + smem = cutlass.memory.SmemAllocator() + red, part = alloc_reduce(smem, NA, WARPS, SPLIT) + cta_reduce(acc, NA, WARPS, SPLIT, warp, lane, rank, red, part) + if warp == 0 and rank == 0: + obase = out.iterator.toint() + for nt in cutlass.range_constexpr(NT): + for i in cutlass.range_constexpr(4): + v = acc[nt * 4 + i] * gsv + if const_expr(RB == 2): + u = acc[(NT + nt) * 4 + i] * gsv + v = silu_mul(round_bf16(v), round_bf16(u)) + row = nb * 16 + g + 8 * (i // 2) + tok = mt * (8 * NT) + nt * 8 + 2 * t + (i % 2) + st_out( + obase + (Int64(tok) * stride_o + row) * 2, + v, + Int32(tok < M), + ) + + return Nvfp4Gemv + + +_kernel_cls: Any = None +_compiled: dict[tuple, Any] = {} + + +def _compile(key: tuple): + global _kernel_cls + import cutlass + import cutlass.cute as cute + from cuda.bindings.driver import CUstream + from quack.compile_utils import make_fake_tensor + + k, nt, cfg, pdl, silu_mul, bf16 = key + if _kernel_cls is None: + _kernel_cls = _build_kernel_class() + sym = cute.sym_int + x = make_fake_tensor(cutlass.BFloat16, (sym(), k), divisibility=8) + w = make_fake_tensor(cutlass.Uint8, (sym(),), divisibility=16) + s = make_fake_tensor(cutlass.Uint8, (sym(),), divisibility=4) + gs = make_fake_tensor(cutlass.Float32, (1,), divisibility=1) + out = make_fake_tensor(cutlass.BFloat16, (sym(), sym()), divisibility=1) + kernel = _kernel_cls(k, nt, cfg.warps, cfg.split, pdl, silu_mul, bf16) + stream = CUstream(torch.cuda.current_stream().cuda_stream) + return cute.compile( + kernel, x, w, s, gs, out, 1, 128, k, 128, stream, options="--enable-tvm-ffi" + ) + + +def _get_compiled(n: int, k: int, nt: int, silu_mul: bool, pdl: bool, bf16: bool): + cfg = _configs(n, k, silu_mul, bf16)[nt] + key = (k, nt, cfg, pdl, silu_mul, bf16) + fn = _compiled.get(key) + if fn is None: + fn = _compiled[key] = _compile(key) + return fn + + +def _launch(x, w, s, gs, n, silu_mul, pdl, bf16) -> torch.Tensor: + from cuda.bindings.driver import CUstream + + m, k = x.shape + assert m > 0 and x.stride(1) == 1 and x.dtype == torch.bfloat16 + nt = pick_nt(n, k, m, silu_mul, bf16) + assert nt is not None, (n, k, m, silu_mul, bf16) + out = torch.empty(m, n // 2 if silu_mul else n, dtype=x.dtype, device=x.device) + fn = _get_compiled(n, k, nt, silu_mul, pdl, bf16) + fn( + x, + w, + s, + gs, + out, + m, + n, + x.stride(0), + out.stride(0), + CUstream(torch.cuda.current_stream().cuda_stream), + ) + return out + + +def nvfp4_gemv( + x: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + global_scale: torch.Tensor, + silu_mul: bool = False, + pdl: bool = True, +) -> torch.Tensor: + """Small-M W4A16 GEMM: ``x @ dequant(weight).T * global_scale``. + + Args: + x: BF16 ``[M, K]`` activations with unit inner stride; ``(N, K)`` + and M must be covered by the tuned tables (``pick_nt``). + weight: uint8 ``[N, K / 2]`` packed E2M1 codes (ModelOpt layout). + weight_scale: 128x4-swizzled E4M3 block scales (``swizzle_blockscale``). + global_scale: FP32 scalar tensor (one element). + silu_mul: Treat the output as ``[gate | up]`` halves and return the + BF16 ``silu(gate) * up`` (``[M, N / 2]``) of the BF16 gate / up + outputs (exact SiLU, one BF16 rounding). + pdl: Launch with programmatic dependent launch: weight loads start + while the previous kernel drains. Disable after kernels that need + the whole GPU to make progress (e.g. attention), whose remaining + CTAs the early, waiting GEMV CTAs would starve. + + Returns: + The BF16 ``[M, N]`` (``[M, N / 2]`` with ``silu_mul``) output. + + """ + assert weight.shape[1] * 2 == x.shape[1] + return _launch( + x, + weight.view(-1), + weight_scale.view(torch.uint8).view(-1), + global_scale.reshape(1), + weight.shape[0], + silu_mul, + pdl, + bf16=False, + ) + + +def bf16_gemv( + x: torch.Tensor, + weight: torch.Tensor, + silu_mul: bool = False, + pdl: bool = True, +) -> torch.Tensor: + """Small-M BF16 GEMM ``x @ weight.T`` (FP32 accumulate); same kernel and + arguments as ``nvfp4_gemv`` with a row-major BF16 ``[N, K]`` weight.""" + assert weight.dtype == torch.bfloat16 and weight.is_contiguous() + assert weight.shape[1] == x.shape[1] + return _launch( + x, + weight.view(torch.uint8).view(-1), + weight.view(torch.uint8).view(-1), + _one(x.device), + weight.shape[0], + silu_mul, + pdl, + bf16=True, + ) + + +_ONES: dict[torch.device, torch.Tensor] = {} + + +def _one(device: torch.device) -> torch.Tensor: + one = _ONES.get(device) + if one is None: + one = _ONES[device] = torch.ones(1, dtype=torch.float32, device=device) + return one + + +def warmup( + n: int, + k: int, + silu_mul: bool = False, + pdl: bool = True, + bf16: bool = False, +) -> None: + """Compile every variant the GEMV may pick for this shape.""" + for nt in _configs(n, k, silu_mul, bf16): + _get_compiled(n, k, nt, silu_mul, pdl, bf16) + + +def _dequant_kernel(): + import triton + import triton.language as tl + + @triton.jit + def kernel( + w_ptr, s_ptr, out_ptr, K: tl.constexpr, BR: tl.constexpr, BK: tl.constexpr + ): + r = tl.program_id(0) * BR + tl.arange(0, BR) + kb = tl.program_id(1) * BK + b = tl.load( + w_ptr + + r[:, None].to(tl.int64) * (K // 2) + + (kb // 2 + tl.arange(0, BK // 2))[None, :] + ) + # fp4 bits -> fp32 bits of e2m1 * 2^-126 (exact), sign included. + lo = ((b & 7).to(tl.int32) << 22) | ((b & 8).to(tl.int32) << 28) + hi = ((b & 0x70).to(tl.int32) << 18) | ((b & 0x80).to(tl.int32) << 24) + v = tl.reshape( + tl.interleave(lo, hi).to(tl.float32, bitcast=True), (BR, BK // 16, 16) + ) + g = kb // 16 + tl.arange(0, BK // 16) + sidx = ((r[:, None] // 128) * (K // 64) + g[None, :] // 4).to(tl.int64) * 512 + sidx += ((r % 32) * 16 + ((r % 128) // 32) * 4)[:, None] + g[None, :] % 4 + sc = tl.load(s_ptr + sidx).to(tl.float32) + v = tl.reshape(v * 8.507059173023462e37 * sc[:, :, None], (BR, BK)) + out = out_ptr + r[:, None].to(tl.int64) * K + (kb + tl.arange(0, BK))[None, :] + tl.store(out, v.to(tl.bfloat16)) + + return kernel + + +_dequant = None + + +def dequantize_rows(weight: torch.Tensor, weight_scale: torch.Tensor) -> torch.Tensor: + """Exact BF16 ``e2m1 * e4m3`` of all weight rows (global scale not + applied); ``weight_scale`` is 128x4-swizzled.""" + global _dequant + if _dequant is None: + _dequant = _dequant_kernel() + rows, k = weight.shape[0], weight.shape[1] * 2 + assert rows % 16 == 0 and k % 256 == 0 + out = torch.empty(rows, k, dtype=torch.bfloat16, device=weight.device) + _dequant[(rows // 16, k // 256)]( + weight, + weight_scale.view(torch.float8_e4m3fn).view(-1), + out, + k, + 16, + 256, + num_warps=4, + ) + return out diff --git a/vllm/models/flint/ops/qk_norm_rope.py b/vllm/models/flint/ops/qk_norm_rope.py new file mode 100644 index 00000000000..0f02e587ed5 --- /dev/null +++ b/vllm/models/flint/ops/qk_norm_rope.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused per-head q/k RMSNorm + NeoX RoPE + paged KV-cache write. + +Numerics follow the unfused vLLM ops it replaces: RMSNorm computes +``bf16(x * rsqrt(mean(x^2) + eps) * w)`` in fp32, RoPE computes +``x * cos - y * sin`` in fp32 from the fp32 cos/sin cache (as the reference +implementation does; vLLM's rotary op would round the cache to bf16 first) and +rounds to bf16, and K/V are stored unchanged in the bf16 cache. +""" + +import torch + +from vllm.triton_utils import tl, triton + +_BLOCK_T = 16 + + +# A fixed token tile and no integer specialization: one compiled variant per +# (RoPE, PDL) pair, all of which CUDA graph capture and warmup reach. +@triton.jit(do_not_specialize=["num_tokens", "num_slots"]) +def _qk_norm_rope_cache_kernel( + qkv_ptr, + qkv_stride, + q_out_ptr, + positions_ptr, + q_weight_ptr, + k_weight_ptr, + cos_sin_ptr, + cos_sin_stride, + slot_mapping_ptr, + num_tokens, + num_slots, + k_cache_ptr, + v_cache_ptr, + stride_block, + stride_slot, + stride_head, + block_size, + eps, + NUM_Q_HEADS: tl.constexpr, + NUM_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + APPLY_ROPE: tl.constexpr, + BLOCK_T: tl.constexpr, + LAUNCH_DEPENDENTS: tl.constexpr, +): + if LAUNCH_DEPENDENTS: + # The PDL successor waits on this grid before reading q / K / V. + tl.extra.cuda.gdc_launch_dependents() + head = tl.program_id(1) + half: tl.constexpr = HEAD_DIM // 2 + tokens = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + offs = tl.arange(0, half)[None, :] + t_mask = (tokens < num_tokens)[:, None] + tok64 = tokens.to(tl.int64)[:, None] + row = qkv_ptr + tok64 * qkv_stride + head * HEAD_DIM + x1 = tl.load(row + offs, mask=t_mask, other=0.0).to(tl.float32) + x2 = tl.load(row + half + offs, mask=t_mask, other=0.0).to(tl.float32) + + is_q = head < NUM_Q_HEADS + is_v = head >= NUM_Q_HEADS + NUM_KV_HEADS + is_k = (head >= NUM_Q_HEADS) & (head < NUM_Q_HEADS + NUM_KV_HEADS) + w_ptr = tl.where(is_q, q_weight_ptr, k_weight_ptr) + w1 = tl.load(w_ptr + offs).to(tl.float32) + w2 = tl.load(w_ptr + half + offs).to(tl.float32) + var = (tl.sum(x1 * x1, axis=1) + tl.sum(x2 * x2, axis=1)) / HEAD_DIM + rs = tl.math.rsqrt(var + eps)[:, None] + n1 = (x1 * rs * w1).to(tl.bfloat16).to(tl.float32) + n2 = (x2 * rs * w2).to(tl.bfloat16).to(tl.float32) + if APPLY_ROPE: + pos = tl.load(positions_ptr + tokens, mask=tokens < num_tokens, other=0) + cs = cos_sin_ptr + pos.to(tl.int64)[:, None] * cos_sin_stride + cos = tl.load(cs + offs, mask=t_mask, other=0.0).to(tl.float32) + sin = tl.load(cs + half + offs, mask=t_mask, other=0.0).to(tl.float32) + r1 = n1 * cos - n2 * sin + r2 = n2 * cos + n1 * sin + n1 = r1 + n2 = r2 + o1 = tl.where(is_v, x1, n1).to(tl.bfloat16) + o2 = tl.where(is_v, x2, n2).to(tl.bfloat16) + + out = q_out_ptr + (tok64 * NUM_Q_HEADS + head) * HEAD_DIM + tl.store(out + offs, o1, mask=t_mask & is_q) + tl.store(out + half + offs, o2, mask=t_mask & is_q) + # Keep the normalized key in qkv as well; attention may read it. + tl.store(row + offs, o1, mask=t_mask & is_k) + tl.store(row + half + offs, o2, mask=t_mask & is_k) + + slot = tl.load(slot_mapping_ptr + tokens, mask=tokens < num_slots, other=-1) + slot = slot.to(tl.int64)[:, None] + kv_head = head - NUM_Q_HEADS - tl.where(is_v, NUM_KV_HEADS, 0) + dst = ( + (slot // block_size) * stride_block + + (slot % block_size) * stride_slot + + kv_head * stride_head + ) + k_mask = (slot >= 0) & is_k + v_mask = (slot >= 0) & is_v + tl.store(k_cache_ptr + dst + offs, o1, mask=k_mask) + tl.store(k_cache_ptr + dst + half + offs, o2, mask=k_mask) + tl.store(v_cache_ptr + dst + offs, o1, mask=v_mask) + tl.store(v_cache_ptr + dst + half + offs, o2, mask=v_mask) + + +def qk_norm_rope_cache( + qkv: torch.Tensor, + positions: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float, + cos_sin_cache: torch.Tensor | None, + k_cache: torch.Tensor | None, + v_cache: torch.Tensor | None, + slot_mapping: torch.Tensor | None, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + launch_dependents: bool = False, +) -> torch.Tensor: + """Normalize and rotate q/k from ``qkv`` and write k/v into the cache. + + Args: + qkv: [T, (Hq + 2 * Hkv) * D] bf16; k is normalized in place. + positions: [T] token positions. + q_weight: [D] q RMSNorm weight. + k_weight: [D] k RMSNorm weight. + eps: RMSNorm epsilon. + cos_sin_cache: [max_pos, D] fp32 (cos | sin), or None for no RoPE. + k_cache: [num_blocks, block_size, Hkv, D] view (any strides with a + contiguous last dim), or None to skip the cache write. + v_cache: Same shape and strides as ``k_cache``. + slot_mapping: [<= T] int64; negative slots are skipped. + num_q_heads: Hq. + num_kv_heads: Hkv. + head_dim: D. + launch_dependents: Let a PDL successor that waits on the grid + dependency before reading q / K / V (Flint decode attention) start + its prologue early. Must stay False before other kernels. + + Returns: + q: [T, Hq * D] bf16. + + """ + num_tokens = qkv.shape[0] + q = torch.empty( + (num_tokens, num_q_heads * head_dim), dtype=qkv.dtype, device=qkv.device + ) + if num_tokens == 0: + return q + if k_cache is not None and v_cache is not None and slot_mapping is not None: + assert k_cache.stride() == v_cache.stride() and k_cache.stride(-1) == 1 + strides = k_cache.stride()[:3] + block_size = k_cache.shape[1] + num_slots = slot_mapping.shape[0] + else: + k_cache = v_cache = slot_mapping = qkv + strides = (0, 0, 0) + block_size = 1 + num_slots = 0 + grid = (triton.cdiv(num_tokens, _BLOCK_T), num_q_heads + 2 * num_kv_heads) + _qk_norm_rope_cache_kernel[grid]( + qkv, + qkv.stride(0), + q, + positions, + q_weight, + k_weight, + qkv if cos_sin_cache is None else cos_sin_cache, + 0 if cos_sin_cache is None else cos_sin_cache.stride(0), + slot_mapping, + num_tokens, + num_slots, + k_cache, + v_cache, + *strides, + block_size, + eps, + NUM_Q_HEADS=num_q_heads, + NUM_KV_HEADS=num_kv_heads, + HEAD_DIM=head_dim, + APPLY_ROPE=cos_sin_cache is not None, + BLOCK_T=_BLOCK_T, + LAUNCH_DEPENDENTS=launch_dependents, + num_warps=4, + ) + return q diff --git a/vllm/models/flint/ops/qkv_fused.py b/vllm/models/flint/ops/qkv_fused.py new file mode 100644 index 00000000000..43b7bdfcda1 --- /dev/null +++ b/vllm/models/flint/ops/qkv_fused.py @@ -0,0 +1,583 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Fused decode qkv for Flint at M <= 8 tokens (CuTe DSL), two kernels: + +1. ``fused_add_rms_norm`` (or ``rms_norm`` for the first layer) writing the + normalized activations ``xn`` and the new residual; it releases its PDL + dependents immediately so the second kernel's weights stream in meanwhile. +2. the qkv GEMV (same fragment math, K split and summation order as + ``nvfp4_gemv`` at M <= 8) whose 16-row CTAs form one cluster per head and + send their BF16 rows to rank 0 through DSMEM, which applies per-head q/k + RMSNorm, NeoX RoPE (FP32 cos/sin table, FP32 math) and the paged KV-cache + write like ``qk_norm_rope_cache``. + +Numerics are those of ``norm.fused_add_rms_norm`` / ``norm.rms_norm`` -> +GEMV -> ``qk_norm_rope_cache``; the per-head q/k norm sums its squares in a +different order, so q/k can differ from that path by one BF16 ulp. +""" + +from functools import cache + +import torch + +from .nvfp4_gemv import _one + +_K = 2048 +_HEAD = 128 +_CHUNK = 64 +_WARPS = 8 +MAX_M = 8 + + +@cache +def _kernels(): + import cutlass + import cutlass.cute as cute + from cutlass import Float32, Int32, Int64, Uint32, const_expr + from cutlass.cutlass_dsl import dsl_user_op + + from .cute_primitives import ( + bf16x2, + decode, + fma, + ld_u8, + ld_w8, + ld_w16, + ld_x16, + lds_f32, + mma, + ptx, + round_bf16, + scale_pair, + smem_addr, + stg_bf16, + stg_u32x4, + sts_f32, + ) + + @dsl_user_op + def ldg128(addr, *, loc=None, ip=None): + return ptx( + "ld.global.v4.b32 {$0,$1,$2,$3}, [$4];", + "=r,=r,=r,=r,l", + [Int64(addr)], + (Uint32,) * 4, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def ldg128_nc(addr, *, loc=None, ip=None): + return ptx( + "ld.global.nc.v4.b32 {$0,$1,$2,$3}, [$4];", + "=r,=r,=r,=r,l", + [Int64(addr)], + (Uint32,) * 4, + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def unpack(w, *, loc=None, ip=None): + return ptx( + "{ .reg .b16 l, h; mov.b32 {l, h}, $2; cvt.f32.bf16 $0, l; " + "cvt.f32.bf16 $1, h; }", + "=f,=f,r", + [Uint32(w)], + (Float32, Float32), + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def rsqrt(x, *, loc=None, ip=None): + return ptx( + "rsqrt.approx.ftz.f32 $0, $1;", + "=f,f", + [Float32(x)], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def ld_bf16(addr, *, loc=None, ip=None): + return ptx( + "{ .reg .b16 h; ld.global.nc.b16 h, [$1]; cvt.f32.bf16 $0, h; }", + "=f,l", + [Int64(addr)], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def ld_f32(addr, *, loc=None, ip=None): + return ptx( + "ld.global.nc.f32 $0, [$1];", + "=f,l", + [Int64(addr)], + Float32, + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def ld_i64(addr, *, loc=None, ip=None): + return ptx( + "ld.global.nc.s64 $0, [$1];", + "=l,l", + [Int64(addr)], + Int64, + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def ld_i64_pred(addr, pred, *, loc=None, ip=None): + return ptx( + "{ .reg .pred p; setp.ne.b32 p, $2, 0; mov.b64 $0, -1; " + "@p ld.global.nc.s64 $0, [$1]; }", + "=l,l,r", + [Int64(addr), Int32(pred)], + Int64, + effects=False, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def st_cluster0_f32(addr, v, *, loc=None, ip=None): + ptx( + "{ .reg .b32 r; mapa.shared::cluster.u32 r, $0, 0; " + "st.shared::cluster.f32 [r], $1; }", + "r,f,~{memory}", + [Uint32(addr), Float32(v)], + loc=loc, + ip=ip, + ) + + class Norm: + """Grid (M,): one CTA of 256 threads per token row, thread i owns + elements 8i .. 8i + 7.""" + + def __init__(self, has_residual): + self.has_residual = has_residual + + @cute.jit + def __call__(self, x, res_in, res_out, xn, norm_w, M: Int32, eps: Float32, + stream): # fmt: skip + self.kernel(x, res_in, res_out, xn, norm_w, eps).launch( + grid=[M, 1, 1], block=[256, 1, 1], stream=stream, use_pdl=True + ) + + @cute.kernel + def kernel(self, x, res_in, res_out, xn, norm_w, eps: Float32): + tidx, _, _ = cute.arch.thread_idx() + tok, _, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + smem = cutlass.memory.SmemAllocator() + var_ptr = smem.allocate_array(Float32, 8, byte_alignment=16) + varb = smem_addr(var_ptr) + nwv = ldg128_nc(norm_w.iterator.toint() + tidx * 16) + cute.arch.griddepcontrol_wait() + # The GEMV only needs our output after its own grid-dependency + # wait: let it launch and stream its weights now. + cute.arch.griddepcontrol_launch_dependents() + off = Int64(tok) * (_K * 2) + tidx * 16 + # s = x (+ residual) in fp32; the variance and the output use the + # unrounded s, the new residual is bf16(s). + hv = ldg128(x.iterator.toint() + off) + sv = [v for i in range(4) for v in unpack(hv[i])] + if const_expr(self.has_residual): + rv = ldg128(res_in.iterator.toint() + off) + rf = [v for i in range(4) for v in unpack(rv[i])] + sv = [sv[j] + rf[j] for j in range(8)] + stg_u32x4( + res_out.iterator.toint() + off, + *[bf16x2(sv[2 * i], sv[2 * i + 1]) for i in range(4)], + ) + acc_s = Float32(0.0) + for i in cutlass.range_constexpr(4): + acc_s = acc_s + sv[2 * i] * sv[2 * i] + sv[2 * i + 1] * sv[2 * i + 1] + for o in cutlass.range_constexpr(5): + acc_s = acc_s + cute.arch.shuffle_sync_bfly(acc_s, 1 << o) + if lane == 0: + sts_f32(varb + warp * 4, acc_s) + cute.arch.sync_threads() + v = Float32(0.0) + for w2 in cutlass.range_constexpr(_WARPS): + v = v + lds_f32(varb + w2 * 4) + rs = rsqrt(v / Float32(_K) + eps) + outw = [] + for i in cutlass.range_constexpr(4): + wlo, whi = unpack(nwv[i]) + outw.append(bf16x2(sv[2 * i] * rs * wlo, sv[2 * i + 1] * rs * whi)) + stg_u32x4(xn.iterator.toint() + off, *outw) + + class QkvRope: + """Grid (N / 16) CTAs in clusters of 8 (one 128-row head each).""" + + def __init__(self, bf16w, rope, num_q_heads, num_kv_heads): + self.bf16w = bf16w + self.rope = rope + self.hq = num_q_heads + self.hkv = num_kv_heads + + @cute.jit + def __call__(self, xn, w, s, gs, qn_w, kn_w, cos_sin, positions, slots, + q_out, k_cache, v_cache, M: Int32, num_slots: Int32, + qk_eps: Float32, stride_block: Int64, stride_slot: Int64, + stride_head: Int64, block_size: Int32, stream): # fmt: skip + n = (self.hq + 2 * self.hkv) * _HEAD + self.kernel( + xn, w, s, gs, qn_w, kn_w, cos_sin, positions, slots, q_out, + k_cache, v_cache, M, num_slots, qk_eps, stride_block, + stride_slot, stride_head, block_size, + ).launch( + grid=[n // 16, 1, 1], + block=[32 * _WARPS, 1, 1], + cluster=[8, 1, 1], + stream=stream, + use_pdl=True, + ) # fmt: skip + + @cute.kernel + def kernel(self, xn, w, s, gs, qn_w, kn_w, cos_sin, positions, slots, + q_out, k_cache, v_cache, M: Int32, num_slots: Int32, + qk_eps: Float32, stride_block: Int64, stride_slot: Int64, + stride_head: Int64, block_size: Int32): # fmt: skip + BF16W: cutlass.Constexpr = self.bf16w + KC: cutlass.Constexpr = _K // _CHUNK + CPW: cutlass.Constexpr = KC // _WARPS + NACC: cutlass.Constexpr = 8 + KB: cutlass.Constexpr = KC * (128 if self.bf16w else 32) + QW: cutlass.Constexpr = 16 if self.bf16w else 4 + tidx, _, _ = cute.arch.thread_idx() + nb, _, _ = cute.arch.block_idx() + head = nb // 8 + rank = nb % 8 + warp = tidx // 32 + lane = tidx % 32 + g = lane // 4 + t = lane % 4 + c_begin = warp * CPW + + smem = cutlass.memory.SmemAllocator() + red_ptr = smem.allocate_array(Float32, _WARPS * 4 * 32, byte_alignment=16) + red = smem_addr(red_ptr) + hb_ptr = smem.allocate_array(Float32, _HEAD * 8, byte_alignment=16) + hb = smem_addr(hb_ptr) + + # ---- weights and epilogue inputs (before the grid dependency) ---- + qs = cute.make_rmem_tensor((CPW * QW,), Uint32) + ss = cute.make_rmem_tensor((CPW * 2,), Uint32) + row_g = nb * 16 + g + row_g8 = row_g + 8 + if const_expr(BF16W): + wbase = w.iterator.toint() + Int64(row_g) * KB + c_begin * 128 + t * 32 + for c in cutlass.range_constexpr(CPW): + for h in cutlass.range_constexpr(4): + off = (h // 2) * 8 * KB + c * 128 + (h % 2) * 16 + wv = ld_w16(wbase + off) + for i in cutlass.range_constexpr(4): + qs[c * QW + h * 4 + i] = wv[i] + else: + wbase = w.iterator.toint() + Int64(row_g) * KB + c_begin * 32 + t * 8 + sbase_g = ( + s.iterator.toint() + + Int64(row_g // 128) * (KC * 512) + + (row_g % 32) * 16 + + ((row_g % 128) // 32) * 4 + + t + + c_begin * 512 + ) + sbase_g8 = ( + s.iterator.toint() + + Int64(row_g8 // 128) * (KC * 512) + + (row_g8 % 32) * 16 + + ((row_g8 % 128) // 32) * 4 + + t + + c_begin * 512 + ) + for c in cutlass.range_constexpr(CPW): + qa = ld_w8(wbase + c * 32) + qb = ld_w8(wbase + 8 * KB + c * 32) + qs[c * QW + 0] = qa[0] + qs[c * QW + 1] = qa[1] + qs[c * QW + 2] = qb[0] + qs[c * QW + 3] = qb[1] + ss[c * 2] = ld_u8(sbase_g + c * 512) + ss[c * 2 + 1] = ld_u8(sbase_g8 + c * 512) + gsv = Float32(1.0) + if const_expr(not BF16W): + gsv = gs[0] + is_q = head < self.hq + is_v = head >= self.hq + self.hkv + wptr = cutlass.select_(is_q, qn_w.iterator.toint(), kn_w.iterator.toint()) + hw = [ld_bf16(Int64(wptr) + (lane + 32 * e) * 2) for e in range(4)] + etc = cutlass.min(warp, M - 1) + pos = Int64(0) + if const_expr(self.rope): + pos = ld_i64(positions.iterator.toint() + Int64(etc) * 8) + slot = ld_i64_pred( + slots.iterator.toint() + Int64(etc) * 8, Int32(etc < num_slots) + ) + cs_v = [Float32(1.0)] * 2 + sn_v = [Float32(0.0)] * 2 + if const_expr(self.rope): + # FP32 cos | sin table rows of 128 floats. + cbase = cos_sin.iterator.toint() + pos * (_HEAD * 4) + for e in cutlass.range_constexpr(2): + cs_v[e] = ld_f32(cbase + (lane + 32 * e) * 4) + sn_v[e] = ld_f32(cbase + (64 + lane + 32 * e) * 4) + cute.arch.cluster_arrive_relaxed() + + cute.arch.griddepcontrol_wait() + xbase = xn.iterator.toint() + (Int64(c_begin) * _CHUNK + t * 16) * 2 + xall = [ + ld_x16(xbase + Int64(g) * (_K * 2) + c * 128, Int32(g < M)) + for c in range(CPW) + ] + + # ---- GEMV ---- + acc = cute.make_rmem_tensor((NACC * 4,), Float32) + acc.fill(0.0) + for c in cutlass.range_constexpr(CPW): + xv = xall[c] + q0 = c * QW + if const_expr(not BF16W): + sg = scale_pair(ss[c * 2]) + sg8 = scale_pair(ss[c * 2 + 1]) + for j in cutlass.range_constexpr(4): + if const_expr(BF16W): + a = [ + qs[q0 + 2 * j], + qs[q0 + 8 + 2 * j], + qs[q0 + 2 * j + 1], + qs[q0 + 8 + 2 * j + 1], + ] + else: + a = decode(qs[q0 + j // 2], qs[q0 + 2 + j // 2], j, sg, sg8) + ab = ((c * 4 + j) % NACC) * 4 + r = mma( + [acc[ab], acc[ab + 1], acc[ab + 2], acc[ab + 3]], + a, + xv[2 * j], + xv[2 * j + 1], + ) + for i in cutlass.range_constexpr(4): + acc[ab + i] = r[i] + for h in cutlass.range_constexpr(1, NACC): + for i in cutlass.range_constexpr(4): + acc[i] = acc[i] + acc[h * 4 + i] + cute.arch.griddepcontrol_launch_dependents() + for i in cutlass.range_constexpr(4): + sts_f32(red + ((warp * 4 + i) * 32 + lane) * 4, acc[i]) + cute.arch.sync_threads() + cute.arch.cluster_wait() + if warp == 0: + for i in cutlass.range_constexpr(4): + v = acc[i] + for w2 in cutlass.range_constexpr(1, _WARPS): + v = v + lds_f32(red + ((w2 * 4 + i) * 32 + lane) * 4) + v = round_bf16(v * gsv) + row = rank * 16 + g + 8 * (i // 2) + tok = 2 * t + (i % 2) + st_cluster0_f32(hb + (row * 8 + tok) * 4, v) + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # ---- per-head epilogue on rank 0: warp = token, lane covers rows + # lane + 32 * r; RoPE pairs (lane, lane + 64), (lane + 32, lane + 96). + if (rank == 0) & (warp < M): + y = [lds_f32(hb + ((lane + 32 * e) * 8 + warp) * 4) for e in range(4)] + sq = y[0] * y[0] + y[1] * y[1] + y[2] * y[2] + y[3] * y[3] + for o in cutlass.range_constexpr(5): + sq = sq + cute.arch.shuffle_sync_bfly(sq, 1 << o) + rs = rsqrt(sq / Float32(_HEAD) + qk_eps) + nv = [round_bf16(y[e] * rs * hw[e]) for e in range(4)] + if const_expr(self.rope): + rot = [None] * 4 + for e in cutlass.range_constexpr(2): + rot[e] = fma(nv[e], cs_v[e], -(nv[e + 2] * sn_v[e])) + rot[e + 2] = fma(nv[e + 2], cs_v[e], nv[e] * sn_v[e]) + nv = rot + out = [cutlass.select_(is_v, y[e], nv[e]) for e in range(4)] + if is_q: + qrow = ( + q_out.iterator.toint() + + (Int64(warp) * (self.hq * _HEAD) + head * _HEAD) * 2 + ) + for e in cutlass.range_constexpr(4): + stg_bf16(qrow + (lane + 32 * e) * 2, out[e]) + elif slot >= 0: + kvh = cutlass.select_( + is_v, head - self.hq - self.hkv, head - self.hq + ) + base = cutlass.select_( + is_v, v_cache.iterator.toint(), k_cache.iterator.toint() + ) + bs = Int64(block_size) + dst = ( + Int64(base) + + ( + (slot // bs) * stride_block + + (slot % bs) * stride_slot + + Int64(kvh) * stride_head + ) + * 2 + ) + for e in cutlass.range_constexpr(4): + stg_bf16(dst + (lane + 32 * e) * 2, out[e]) + + return Norm, QkvRope + + +@cache +def _compile_norm(has_residual): + import cutlass + import cutlass.cute as cute + from cutlass import Float32, Int32 + from quack.compile_utils import make_fake_tensor + + sym = cute.sym_int + bf16 = cutlass.BFloat16 + args = [make_fake_tensor(bf16, (sym(), _K), divisibility=8) for _ in range(4)] + args.append(make_fake_tensor(bf16, (_K,), divisibility=8)) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + _kernels()[0](has_residual), *args, Int32(1), Float32(1e-6), stream, + options="--enable-tvm-ffi", + ) # fmt: skip + + +@cache +def _compile_qkv(bf16w, rope, num_q_heads, num_kv_heads): + import cutlass + import cutlass.cute as cute + from cutlass import Float32, Int32, Int64 + from quack.compile_utils import make_fake_tensor + + sym = cute.sym_int + bf16 = cutlass.BFloat16 + args = [ + make_fake_tensor(bf16, (sym(), _K), divisibility=8), # xn + make_fake_tensor(cutlass.Uint8, (sym(),), divisibility=16), # weight + make_fake_tensor(cutlass.Uint8, (sym(),), divisibility=4), # scales + make_fake_tensor(Float32, (1,), divisibility=1), # global scale + make_fake_tensor(bf16, (_HEAD,), divisibility=8), # q norm w + make_fake_tensor(bf16, (_HEAD,), divisibility=8), # k norm w + make_fake_tensor(Float32, (sym(), sym()), divisibility=4), # cos | sin + make_fake_tensor(cutlass.Int64, (sym(),), divisibility=1), # positions + make_fake_tensor(cutlass.Int64, (sym(),), divisibility=1), # slots + make_fake_tensor(bf16, (sym(), num_q_heads * _HEAD), divisibility=8), + make_fake_tensor(bf16, (sym(),), divisibility=8), # k cache (flat) + make_fake_tensor(bf16, (sym(),), divisibility=8), # v cache (flat) + ] + kernel = _kernels()[1](bf16w, rope, num_q_heads, num_kv_heads) + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + return cute.compile( + kernel, *args, Int32(1), Int32(1), Float32(1e-6), Int64(1), Int64(1), + Int64(1), Int32(16), stream, options="--enable-tvm-ffi", + ) # fmt: skip + + +def flint_qkv_fused( + x: torch.Tensor, + residual: torch.Tensor | None, + norm_weight: torch.Tensor, + eps: float, + weight: torch.Tensor, + weight_scale: torch.Tensor | None, + global_scale: torch.Tensor | None, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + qk_eps: float, + cos_sin_cache: torch.Tensor | None, + positions: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + slot_mapping: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Residual-add RMSNorm + qkv GEMV + q/k norm, RoPE, cache write for + ``M <= 8`` tokens. + + Args: + x: [M, 2048] bf16 attention input (contiguous). + residual: [M, 2048] bf16 residual, or None for the first layer. + norm_weight: Input RMSNorm weight. + eps: Input RMSNorm epsilon. + weight: qkv weight, NVFP4 (uint8 [N, 1024]) or BF16 [N, 2048]. + weight_scale: NVFP4 128x4-swizzled block scales (None for BF16). + global_scale: NVFP4 fp32 global scale (None for BF16). + q_norm_weight: Per-head q RMSNorm weight. + k_norm_weight: Per-head k RMSNorm weight. + qk_eps: q/k RMSNorm epsilon. + cos_sin_cache: [max_pos, 128] fp32 (cos | sin), or None (no RoPE). + positions: [M] int64. + k_cache: [blocks, block_size, kv_heads, 128] view. + v_cache: Same shape and strides as ``k_cache``. + slot_mapping: [<= M] int64; negative slots and tokens beyond it are + not written. + num_q_heads: Number of q heads. + num_kv_heads: Number of KV heads. + + Returns: + (q [M, num_q_heads * 128] bf16, residual_out [M, 2048] bf16). + + """ + m = x.shape[0] + assert 0 < m <= MAX_M and x.shape[1] == _K and x.is_contiguous() + bf16w = weight.dtype == torch.bfloat16 + n = (num_q_heads + 2 * num_kv_heads) * _HEAD + assert weight.shape[0] == n and weight.is_contiguous() + assert k_cache.stride() == v_cache.stride() and k_cache.stride(-1) == 1 + has_residual = residual is not None + if residual is not None: + assert residual.is_contiguous() + res_out = torch.empty_like(x) + else: + residual = res_out = x + xn = torch.empty_like(x) + _compile_norm(has_residual)(x, residual, res_out, xn, norm_weight, m, eps) + q = torch.empty(m, num_q_heads * _HEAD, dtype=x.dtype, device=x.device) + w = weight.view(torch.uint8).view(-1) + s = w if weight_scale is None else weight_scale.view(torch.uint8).view(-1) + gs = global_scale.reshape(1) if global_scale is not None else _one(x.device) + rope = cos_sin_cache is not None + if cos_sin_cache is not None: + assert cos_sin_cache.dtype == torch.float32 and cos_sin_cache.shape[1] == _HEAD + cs = cos_sin_cache + else: + cs = no_rope_table(x.device) + _compile_qkv(bf16w, rope, num_q_heads, num_kv_heads)( + xn, w, s, gs, q_norm_weight, k_norm_weight, cs, positions, slot_mapping, + q, _base(k_cache), _base(v_cache), m, slot_mapping.shape[0], qk_eps, + k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), k_cache.shape[1], + ) # fmt: skip + return q, res_out + + +def _base(t: torch.Tensor) -> torch.Tensor: + return torch.as_strided(t, (_HEAD,), (1,)) + + +_NO_ROPE: dict[torch.device, torch.Tensor] = {} + + +def no_rope_table(device: torch.device) -> torch.Tensor: + """Placeholder fp32 cos|sin argument for kernels compiled without RoPE.""" + t = _NO_ROPE.get(device) + if t is None: + t = _NO_ROPE[device] = torch.zeros(1, _HEAD, device=device) + return t diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index e38b6c4c1fa..c5c9438ac56 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -169,6 +169,7 @@ _CONFIG_REGISTRY: dict[str, type[PreTrainedConfig]] = LazyConfigDict( **{"unlimited-ocr": "UnlimitedOCRConfig"}, unlimited_ocr="UnlimitedOCRConfig", **{"deepseek_v41": "DeepseekV41Config"}, + flint="FlintConfig", inkling_mm_model="InklingMMConfig", inkling_model="InklingModelConfig", ) @@ -821,6 +822,10 @@ def maybe_override_with_speculators( Tuple of (resolved_model, resolved_tokenizer, speculative_config) """ + # A single local file (e.g. a .gguf checkpoint) has no speculators config. + if os.path.isfile(model): + return model, tokenizer, vllm_speculative_config + kwargs["local_files_only"] = huggingface_hub.constants.HF_HUB_OFFLINE config_dict, _ = PreTrainedConfig.get_config_dict( model, diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 538cf5e9cd2..188e39c115c 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -74,6 +74,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "Step5Config": "vllm.transformers_utils.configs.step5", "Step5VLConfig": "vllm.transformers_utils.configs.step5", "Qwen3ASRConfig": "vllm.transformers_utils.configs.qwen3_asr", + "FlintConfig": "vllm.models.flint.configs", "InklingModelConfig": "vllm.models.inkling.configs", "InklingAudioConfig": "vllm.models.inkling.configs", "InklingVisionConfig": "vllm.models.inkling.configs", @@ -144,6 +145,7 @@ __all__ = [ "Step5Config", "Step5VLConfig", "Qwen3ASRConfig", + "FlintConfig", "InklingModelConfig", "InklingAudioConfig", "InklingVisionConfig",