vllm.models.flint.model

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.

Classes:

• FlintAttention –
• FlintSparseMoeBlock –

FlintAttention

Bases: Module

Methods:

• forward_ll –

  Small pure-decode batches: two fused kernels for the residual-add

Source code in vllm/models/flint/model.py

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

_cache_views()

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

Source code in vllm/models/flint/model.py

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)

_decode_plan()

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

Source code in vllm/models/flint/model.py

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

forward_ll(positions, hidden_states, residual, input_layernorm)

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.

Source code in vllm/models/flint/model.py

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

FlintSparseMoeBlock

Bases: Module

Methods:

• forward_ll –

  (mlp(norm(attn_out + residual)), attn_out + residual) via the

Source code in vllm/models/flint/model.py

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)

forward_ll(attn_out, residual, norm)

(mlp(norm(attn_out + residual)), attn_out + residual) via the fused decode MoE block (NVFP4, M <= LL_MOE_MAX_TOKENS), else None.

Source code in vllm/models/flint/model.py

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

_flint_kernels(quant_config, parallel=None)

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.

Source code in vllm/models/flint/model.py

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
    )