[wanfengcxz]

No description provided.

[silencelamb]

[P1] MTP short prefill corrupts the convolution ring state

While reviewing this PR (40cfc6c) together with InternLM/lmdeploy#4999, I reproduced a remaining correctness issue in causal_conv1d_fn's ring write. This helper is unchanged by this PR, but the new LMDeploy adapter still calls it.

The write indices are calculated from each sequence's end in the flattened batch and only clamped to global index 0. For a sequence/chunk shorter than the convolution width, the gathered values can come from the previous request, or repeat the first current token instead of retaining the cached history. With width 4 and a 1–2 token prefill, this changes subsequent decode outputs. It also affects an ordinary long prompt whose final prefill chunk is short.

Minimal CPU reproduction

Run the following from a dlinfer checkout at 40cfc6c992a3934c4cf1f747e2748fd201478364 with PyTorch installed. The AST loader executes the two original pure-PyTorch functions without importing the NPU/Triton package initialization. The next-token check uses the original convolution reference with the history that decode would read from the ring.

import ast
from pathlib import Path
from typing import Any, Optional

import torch
import torch.nn.functional as F

# Run from a dlinfer checkout at 40cfc6c. Load the two unchanged CPU
# functions directly, bypassing torch_npu/Triton package initialization.
path = Path('dlinfer/vendor/ascend/triton_ops/causal_conv1d.py')
tree = ast.parse(path.read_text())
tree.body = [n for n in tree.body if isinstance(n, ast.FunctionDef)
             and n.name in ('causal_conv1d_ref', 'causal_conv1d_fn')]
ns = dict(torch=torch, F=F, Any=Any, Optional=Optional, PAD_SLOT_ID=-1)
exec(compile(tree, str(path), 'exec'), ns)
prefill, reference = ns['causal_conv1d_fn'], ns['causal_conv1d_ref']
weight = torch.ones(1, 4)  # one channel, convolution width 4

def run(cache, tokens, lengths, history, ids):
    q = torch.tensor(lengths, dtype=torch.int32)
    h = torch.tensor(history, dtype=torch.int32)
    prefill(
        torch.tensor([tokens], dtype=torch.float32), weight,
        conv_states=cache, has_initial_state=h > 0,
        cache_indices=torch.tensor(ids, dtype=torch.int32),
        query_start_loc=F.pad(q.cumsum(0, dtype=torch.int32), (1, 0)),
        read_conv_offsets=(h[:, None] + torch.arange(-3, 0)) % 6,
        write_conv_offsets=((h + q)[:, None] + torch.arange(-4, 0)) % 6,
    )

def next_output(cache, slot, history):
    offsets = (history + torch.arange(-3, 0)) % 6
    initial = cache[slot].index_select(-1, offsets)
    out, _ = reference(torch.tensor([[2.]]), weight, initial_states=initial)
    return out.item()

# Fresh requests: [10, 20, 30, 40] and [1]. Ring length = 4 + 2.
cache = torch.zeros(3, 1, 6)
run(cache, [10, 20, 30, 40, 1], [4, 1], [0, 0], [1, 2])
expected = torch.zeros_like(cache)
expected[2, 0, 0] = 1
print('fresh cache:', cache[2, 0].tolist())
print('expected:   ', expected[2, 0].tolist())
print('next output:', next_output(cache, 2, 1), 'expected:', next_output(expected, 2, 1))

# A one-token continuation after history [10, 20, 30, 40].
cache = torch.zeros(2, 1, 6)
cache[1, 0, :4] = torch.tensor([10., 20., 30., 40.])
expected = cache.clone()
expected[1, 0, 4] = 1
run(cache, [1], [1], [4], [1])
print('tail cache: ', cache[1, 0].tolist())
print('expected:   ', expected[1, 0].tolist())
print('next output:', next_output(cache, 1, 5), 'expected:', next_output(expected, 1, 5))

Observed output:

fresh cache: [1.0, 0.0, 0.0, 20.0, 30.0, 40.0]
expected:    [1.0, 0.0, 0.0, 0.0, 0.0, 0.0]
next output: 73.0 expected: 2.857722520828247

tail cache:  [10.0, 1.0, 1.0, 1.0, 1.0, 0.0]
expected:    [10.0, 20.0, 30.0, 40.0, 1.0, 0.0]
next output: 4.966535568237305 expected: 73.0

In the first case, request 2's history contains 20/30/40 from request 1. In the second, the one-token continuation overwrites valid history with repeated 1 values.

Suggested fix

For each valid sequence, gather its existing width - 1 history using read_conv_offsets (or use zeros when has_initial_state=False), concatenate that history with the sequence's own new input, and scatter the last width values to its write_conv_offsets. Gather the history before overwriting the ring. This handles short chunks without crossing request boundaries or destroying prior history; merely clamping to the sequence start would still repeat the first new token.

Please add regression tests for fresh and continuing sequences with query lengths 1, 2, width - 1, and >= width, including mixed-length batches and ring wraparound. Check both the resulting cache and subsequent outputs, and compare batched vs. independent requests and chunked vs. one-shot prefill.

Validation here is CPU execution of the original prefill/reference functions. The same failure was also observed through LMDeploy's new metadata/conv adapter with the vendor decode kernel under Triton's CPU interpreter; this is not an end-to-end Ascend NPU test.

[wanfengcxz]

[P1] MTP short prefill corrupts the convolution ring state

While reviewing this PR (40cfc6c) together with InternLM/lmdeploy#4999, I reproduced a remaining correctness issue in causal_conv1d_fn's ring write. This helper is unchanged by this PR, but the new LMDeploy adapter still calls it.

The write indices are calculated from each sequence's end in the flattened batch and only clamped to global index 0. For a sequence/chunk shorter than the convolution width, the gathered values can come from the previous request, or repeat the first current token instead of retaining the cached history. With width 4 and a 1–2 token prefill, this changes subsequent decode outputs. It also affects an ordinary long prompt whose final prefill chunk is short.

Minimal CPU reproduction

Run the following from a dlinfer checkout at 40cfc6c992a3934c4cf1f747e2748fd201478364 with PyTorch installed. The AST loader executes the two original pure-PyTorch functions without importing the NPU/Triton package initialization. The next-token check uses the original convolution reference with the history that decode would read from the ring.

import ast
from pathlib import Path
from typing import Any, Optional

import torch
import torch.nn.functional as F

# Run from a dlinfer checkout at 40cfc6c. Load the two unchanged CPU
# functions directly, bypassing torch_npu/Triton package initialization.
path = Path('dlinfer/vendor/ascend/triton_ops/causal_conv1d.py')
tree = ast.parse(path.read_text())
tree.body = [n for n in tree.body if isinstance(n, ast.FunctionDef)
             and n.name in ('causal_conv1d_ref', 'causal_conv1d_fn')]
ns = dict(torch=torch, F=F, Any=Any, Optional=Optional, PAD_SLOT_ID=-1)
exec(compile(tree, str(path), 'exec'), ns)
prefill, reference = ns['causal_conv1d_fn'], ns['causal_conv1d_ref']
weight = torch.ones(1, 4)  # one channel, convolution width 4

def run(cache, tokens, lengths, history, ids):
    q = torch.tensor(lengths, dtype=torch.int32)
    h = torch.tensor(history, dtype=torch.int32)
    prefill(
        torch.tensor([tokens], dtype=torch.float32), weight,
        conv_states=cache, has_initial_state=h > 0,
        cache_indices=torch.tensor(ids, dtype=torch.int32),
        query_start_loc=F.pad(q.cumsum(0, dtype=torch.int32), (1, 0)),
        read_conv_offsets=(h[:, None] + torch.arange(-3, 0)) % 6,
        write_conv_offsets=((h + q)[:, None] + torch.arange(-4, 0)) % 6,
    )

def next_output(cache, slot, history):
    offsets = (history + torch.arange(-3, 0)) % 6
    initial = cache[slot].index_select(-1, offsets)
    out, _ = reference(torch.tensor([[2.]]), weight, initial_states=initial)
    return out.item()

# Fresh requests: [10, 20, 30, 40] and [1]. Ring length = 4 + 2.
cache = torch.zeros(3, 1, 6)
run(cache, [10, 20, 30, 40, 1], [4, 1], [0, 0], [1, 2])
expected = torch.zeros_like(cache)
expected[2, 0, 0] = 1
print('fresh cache:', cache[2, 0].tolist())
print('expected:   ', expected[2, 0].tolist())
print('next output:', next_output(cache, 2, 1), 'expected:', next_output(expected, 2, 1))

# A one-token continuation after history [10, 20, 30, 40].
cache = torch.zeros(2, 1, 6)
cache[1, 0, :4] = torch.tensor([10., 20., 30., 40.])
expected = cache.clone()
expected[1, 0, 4] = 1
run(cache, [1], [1], [4], [1])
print('tail cache: ', cache[1, 0].tolist())
print('expected:   ', expected[1, 0].tolist())
print('next output:', next_output(cache, 1, 5), 'expected:', next_output(expected, 1, 5))

Observed output:

fresh cache: [1.0, 0.0, 0.0, 20.0, 30.0, 40.0]
expected:    [1.0, 0.0, 0.0, 0.0, 0.0, 0.0]
next output: 73.0 expected: 2.857722520828247

tail cache:  [10.0, 1.0, 1.0, 1.0, 1.0, 0.0]
expected:    [10.0, 20.0, 30.0, 40.0, 1.0, 0.0]
next output: 4.966535568237305 expected: 73.0

In the first case, request 2's history contains 20/30/40 from request 1. In the second, the one-token continuation overwrites valid history with repeated 1 values.

Suggested fix

For each valid sequence, gather its existing width - 1 history using read_conv_offsets (or use zeros when has_initial_state=False), concatenate that history with the sequence's own new input, and scatter the last width values to its write_conv_offsets. Gather the history before overwriting the ring. This handles short chunks without crossing request boundaries or destroying prior history; merely clamping to the sequence start would still repeat the first new token.

Please add regression tests for fresh and continuing sequences with query lengths 1, 2, width - 1, and >= width, including mixed-length batches and ring wraparound. Check both the resulting cache and subsequent outputs, and compare batched vs. independent requests and chunked vs. one-shot prefill.

Validation here is CPU execution of the original prefill/reference functions. The same failure was also observed through LMDeploy's new metadata/conv adapter with the vendor decode kernel under Triton's CPU interpreter; this is not an end-to-end Ascend NPU test.

Confirmed — excellent catch, and the analysis was spot-on. I reproduced your CPU repro exactly (both the cross-request case and the one-token continuation), and traced the root cause to the same lines: the
vectorised ring write gathers each sequence's trailing width positions from the flattened batch x with clamp_min(0) only guarding the global index, so for q < width it either crosses into the previous
request's tokens or repeats the first token, overwriting the history that only lives in the ring.

Fixed in ce25158. The write now follows exactly the semantics you suggested:

• gather the existing width - 1 history via read_conv_offsets before any scatter (zeros where has_initial_state=False),
• concatenate the sequence's own new tokens,
• scatter the trailing width values of that concatenation to write_conv_offsets.

While fixing this I also found the old block scattered rows the compute loop skips (PAD_SLOT_ID rows landed on conv_states[-1], and zero-length rows wrote garbage); the new path skips both.

Regression tests (tests/test_ascend_causal_conv1d_ring.py, 16 cases, runs on CPU — the prefill helper is pure PyTorch): fresh and continuing sequences with q ∈ {1, 2, 3, 4, 5, 8} (covers 1, 2, width−1,
≥width), ring wraparound, mixed-length batches, and chunked-vs-one-shot equivalence. Each case asserts both the resulting cache window and the next-token output against causal_conv1d_ref on the true history,
which also gives the batched-vs-independent check you asked for.

End-to-end on A3 (Qwen3.5-9B, MTP num_spec_tokens=4, greedy, identical batch shapes; only the dlinfer file swapped between runs):

• Mixed batch batch_size=2 num_tokens=342 (341-token request followed by a fresh 1-token request): outputs of the short request diverge at token 13 (buggy) vs (fixed); the long request's own output is
  bit-identical across both builds, as expected since it is the contamination source, not the victim.
• Chunked prefill 349 + 1-token tail chunk (is_chunk=True confirmed in logs): identical for the first 100 tokens, then diverges — consistent with the corruption touching only 2 of ~350 context tokens through
  one layer's conv state, i.e. subtle and easy to misattribute to sampling noise in production.

Your validation caveat (CPU / Triton-interpreter only) was fair — the NPU run above closes that gap. Thanks again for the detailed report; the boundary-crossing gather would have been very hard to spot from
the serving side.

[silencelamb]

LGTM