[Model] Add Recursive Flint (#62938)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

.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

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)

docker/versions.json

@@ -84,6 +84,9 @@
     },
     "RUNAI_MODEL_STREAMER_VERSION": {
       "default": ">=0.15.7"
+    },
+    "GGUF_VERSION": {
+      "default": ">=0.19.0"
     }
   }
 }

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. | ✅︎ | ✅︎ |

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 <tool_call>/<arg_key>/<arg_value>); 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 <think>…</think> 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 `<think>` or `<tool_call>` 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)

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

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)

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)

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 <org>/<repo>` 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

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)

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)

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(

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)

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"),

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}")

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"]

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 ``<name>.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]

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

vllm/models/flint/ops/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

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

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

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()

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(),
+    )

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)

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

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

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

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,

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",