Skip to content

Commit f2c5ec1

Browse files
committed
fix mtp.py typo
Signed-off-by: liugaoji <[email protected]>
1 parent 15ceba8 commit f2c5ec1

File tree

1 file changed

+3
-3
lines changed
  • tensorrt_llm/_torch/speculative

1 file changed

+3
-3
lines changed

tensorrt_llm/_torch/speculative/mtp.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -406,20 +406,20 @@ def forward(
406406
- KV cache: (ABCD) + EFGH (H's KV cache is invalid)
407407
- hidden states: H_E, H_F, H_G, H_H (H_H is invalid)
408408
Draft model:
409-
MPT1:
409+
MTP1:
410410
# For generation request, `mtp_num_modules` of tokens will be used as input.
411411
- input tokens: FGX
412412
- input hidden states: H_E, H_F, H_G
413413
- KV cache: (BCDE) + FGX
414414
- output hidden states: h_F, h_G, h_X
415415
- output next draft token: N
416-
MPT2:
416+
MTP2:
417417
- input tokens: GXN
418418
- input hidden states: H_F, H_G, h_X
419419
- KV cache: (CDEF) + GXN
420420
- output hidden states: h_G, h_X, h_N
421421
- output next draft token: O
422-
MPT3:
422+
MTP3:
423423
- input tokens: XNO
424424
- input hidden states: H_G, H_X, h_N
425425
- KV cache: (DEFG) + XNO

0 commit comments

Comments
 (0)