Restore the TPC causal_conv1d update path, switchable by env var

docs/configuration/env_variables.md

@@ -51,6 +51,7 @@ To enter developer mode use `VLLM_DEVELOPER_MODE`:
 | `VLLM_HANDLE_TOPK_DUPLICATES` | Handles duplicates outside top-k.                                                                                                                                                             | `false`       |
 | `VLLM_CONFIG_HIDDEN_LAYERS`   | Sets the number of hidden layers to run per HPUGraph for model splitting among hidden layers when TP is 1. It improves throughput by reducing inter-token latency limitations in some models. | `1`           |
 | `VLLM_WORKER_MULTIPROC_METHOD` | Sets the Python `multiprocessing` start method used by the `mp` distributed executor backend when launching worker processes. The upstream default is `fork`. On HPU, it is automatically overridden to `spawn` with a warning because forked child processes inherit HPU driver state and can hang on exit. The override is applied when `--distributed-executor-backend` is `mp` or `uni`. With `uni`, no subprocess is created, so the value has no practical effect. With `external_launcher` and `ray`, workers are not started through Python `multiprocessing`, so the value is irrelevant. Set `VLLM_WORKER_MULTIPROC_METHOD=spawn` explicitly to suppress the auto-override warning, or set it to `fork` to opt out of the override, which is not recommended. | `spawn` on HPU (auto-overridden from upstream `fork`) |
+| `VLLM_HPU_CONV1D_DISABLE_TPC` | Runs the PyTorch reference implementation of the hybrid-model `causal_conv1d` decode update instead of the TPC kernel. Provided so both implementations can be compared from one build; leave unset for normal use. | `false`       |
 
 ## Heterogeneous KV Transfer (NIXL)
 

tests/unit_tests/ops/test_causal_conv1d_update_dispatch.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Which implementation hpu_causal_conv1d_update() dispatches to.
+
+`hpu_causal_conv1d_update` has two implementations: the TPC kernel
+`torch.ops.hpu.causal_conv1d_update`, and the PyTorch reference reached through
+`hpu_causal_conv1d_fn_update`. The choice depends on whether the op is present
+and on `VLLM_HPU_CONV1D_DISABLE_TPC`.
+
+These tests cover the choice itself, not the numerics — both sides are stubbed,
+so they run on CPU. Numerical equivalence of the TPC path is covered by
+test_depthwise_conv1d_tpc.py, which needs hardware.
+
+The case that matters most is the unset variable: the TPC kernel is the default,
+so a missing or empty variable must not quietly fall back to the reference.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import vllm_gaudi.ops.causal_conv1d_pytorch as conv1d
+
+VAR = "VLLM_HPU_CONV1D_DISABLE_TPC"
+
+DIM, WIDTH, BATCH = 4, 4, 2
+STATE_LEN = WIDTH - 1
+
+
+@pytest.fixture
+def stubs(monkeypatch):
+    """Replace both implementations with recorders and return the record."""
+    taken: list[str] = []
+
+    def fake_tpc(x_3d, conv_state, weight, bias, activation=False, pad_slot_id=-1):
+        taken.append("tpc")
+        # (out, conv_state_out) with the shapes the caller expects back
+        return torch.zeros_like(x_3d), torch.zeros_like(conv_state)
+
+    def fake_reference(flat_x, *args, **kwargs):
+        taken.append("reference")
+        # The caller reshapes this back to x's layout, so echo the flattened shape
+        # it was handed rather than inventing one.
+        return torch.zeros_like(flat_x)
+
+    # torch.ops.hpu is absent without habana_frameworks, so create it; raising=False
+    # covers both cases.
+    monkeypatch.setattr(torch.ops, "hpu", SimpleNamespace(causal_conv1d_update=fake_tpc), raising=False)
+    monkeypatch.setattr(conv1d, "hpu_causal_conv1d_fn_update", fake_reference)
+    return taken
+
+
+def _call():
+    x = torch.zeros(BATCH, DIM)
+    conv_state = torch.zeros(BATCH, STATE_LEN, DIM)
+    weight = torch.zeros(DIM, WIDTH)
+    bias = torch.zeros(DIM)
+    conv1d.hpu_causal_conv1d_update(
+        x,
+        conv_state,
+        weight,
+        bias,
+        activation="silu",
+        conv_state_indices=torch.arange(BATCH, dtype=torch.int32),
+        query_start_loc=torch.arange(BATCH + 1, dtype=torch.int64),
+    )
+
+
+@pytest.mark.parametrize(
+    "value, expected",
+    [
+        (None, "tpc"),  # unset — the release default
+        ("", "tpc"),
+        ("0", "tpc"),
+        ("false", "tpc"),
+        ("FALSE", "tpc"),
+        ("1", "reference"),
+        ("true", "reference"),
+        ("TRUE", "reference"),
+    ],
+)
+def test_env_var_selects_implementation(monkeypatch, stubs, value, expected):
+    if value is None:
+        monkeypatch.delenv(VAR, raising=False)
+    else:
+        monkeypatch.setenv(VAR, value)
+    _call()
+    assert stubs == [expected]
+
+
+def test_reference_used_when_op_unavailable(monkeypatch, stubs):
+    """Without the TPC op the reference runs, whatever the variable says."""
+    monkeypatch.setattr(torch.ops, "hpu", SimpleNamespace(), raising=False)
+    monkeypatch.delenv(VAR, raising=False)
+    _call()
+    assert stubs == ["reference"]
+
+
+def test_variable_read_per_call(monkeypatch, stubs):
+    """The value is not cached at import, so it can differ between calls."""
+    monkeypatch.delenv(VAR, raising=False)
+    _call()
+    monkeypatch.setenv(VAR, "1")
+    _call()
+    monkeypatch.delenv(VAR, raising=False)
+    _call()
+    assert stubs == ["tpc", "reference", "tpc"]

vllm_gaudi/ops/causal_conv1d_pytorch.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import Callable
 
@@ -375,21 +376,82 @@ def hpu_causal_conv1d_update(
     activation = _normalize_activation(activation)
     dim = weight.size(0)
 
-    flat_x, qsl, reshape_spec = _flatten_inputs_for_update(x, query_start_loc, dim)
-    result = hpu_causal_conv1d_fn_update(
-        flat_x,
-        weight,
-        bias,
-        conv_state,
-        qsl,
-        cache_indices=conv_state_indices,
-        has_initial_state=None,
-        activation=activation,
-        metadata=None,
-        validate_data=validate_data,
-        is_prompt=False,
-    )
-    return reshape_spec.reshape_fn(result)
+    # Use the TPC kernel when it is available, unless it has been switched off.
+    # Read per call, not cached at import, so the setting reaches every worker
+    # process regardless of how it was spawned.
+    disable_tpc = os.environ.get("VLLM_HPU_CONV1D_DISABLE_TPC", "0").lower() in ("1", "true")
+
+    if not disable_tpc and hasattr(torch.ops.hpu, "causal_conv1d_update"):
+        activation_bool = activation in ("silu", "swish")
+        work_dtype = conv_state.dtype
+        # weight: (dim, width) -> (width, dim)
+        weight_for_tpc = weight.t().contiguous().to(work_dtype)
+        bias_for_tpc = bias.to(work_dtype) if bias is not None else None
+
+        # Gather per-batch conv_state using indices
+        if conv_state_indices is not None:
+            num_conv_slots = conv_state.shape[0]
+            safe_idx = torch.remainder(conv_state_indices.long(), num_conv_slots)
+            batch_conv_state = conv_state[safe_idx]
+            batch = safe_idx.numel()
+        else:
+            batch_conv_state = conv_state
+            batch = conv_state.shape[0]
+
+        # Reshape x to (batch, seqlen, dim) for TPC kernel
+        original_x = x
+        if x.dim() == 2:
+            if x.size(1) == dim:
+                x_3d = x.reshape(batch, -1, dim).to(work_dtype)
+            elif x.size(0) == dim:
+                x_3d = x.t().contiguous().reshape(batch, -1, dim).to(work_dtype)
+            else:
+                raise ValueError("Cannot reshape 2-D x for TPC causal_conv1d_update.")
+        elif x.dim() == 3:
+            x_3d = x.to(work_dtype)
+        else:
+            raise ValueError("Unsupported x dimensions for TPC causal_conv1d_update.")
+
+        out, conv_state_out = torch.ops.hpu.causal_conv1d_update(
+            x_3d,
+            batch_conv_state,
+            weight_for_tpc,
+            bias_for_tpc,
+            activation=activation_bool,
+            pad_slot_id=pad_slot_id,
+        )
+
+        # Write back updated conv_state
+        with torch.no_grad():
+            if conv_state_indices is not None:
+                conv_state[safe_idx] = conv_state_out
+            else:
+                conv_state.copy_(conv_state_out)
+
+        # Reshape output to match original x shape
+        if original_x.dim() == 2:
+            if original_x.size(1) == dim:
+                return out.reshape(original_x.shape).to(original_x.dtype)
+            elif original_x.size(0) == dim:
+                return out.reshape(-1, dim).t().contiguous().to(original_x.dtype)
+        return out.to(original_x.dtype)
+
+    else:
+        flat_x, qsl, reshape_spec = _flatten_inputs_for_update(x, query_start_loc, dim)
+        result = hpu_causal_conv1d_fn_update(
+            flat_x,
+            weight,
+            bias,
+            conv_state,
+            qsl,
+            cache_indices=conv_state_indices,
+            has_initial_state=None,
+            activation=activation,
+            metadata=None,
+            validate_data=validate_data,
+            is_prompt=False,
+        )
+        return reshape_spec.reshape_fn(result)
 
 
 def hpu_causal_conv1d_fn_update(