[refactor] init_forward_metadata 3-method ABC + side-channel removal + ForwardMetadata type rename (#26735)
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.7
parent
6c69756fa8
commit
99da43b900
@@ -1,5 +1,6 @@
|
||||
import math
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -1030,15 +1031,15 @@ class TestTRTLLMMLA(CustomTestCase):
|
||||
)
|
||||
req_pool_indices = torch.arange(batch_size, device=config["device"])
|
||||
|
||||
backend.init_forward_metadata_capture_cuda_graph(
|
||||
bs=batch_size,
|
||||
num_tokens=batch_size,
|
||||
capture_fb = SimpleNamespace(
|
||||
batch_size=batch_size,
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=seq_lens,
|
||||
encoder_lens=None,
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
positions=torch.arange(batch_size, device=config["device"]),
|
||||
spec_info=None,
|
||||
)
|
||||
backend.init_forward_metadata_out_graph(capture_fb, in_capture=True)
|
||||
|
||||
# Verify capture metadata
|
||||
self.assertIn(batch_size, backend.decode_cuda_graph_metadata)
|
||||
@@ -1054,16 +1055,15 @@ class TestTRTLLMMLA(CustomTestCase):
|
||||
device=config["device"],
|
||||
)
|
||||
|
||||
backend.init_forward_metadata_replay_cuda_graph(
|
||||
bs=batch_size,
|
||||
replay_fb = SimpleNamespace(
|
||||
batch_size=batch_size,
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
req_pool_indices=req_pool_indices,
|
||||
seq_lens=new_seq_lens,
|
||||
seq_lens_sum=new_seq_lens.sum().item(),
|
||||
encoder_lens=None,
|
||||
forward_mode=ForwardMode.DECODE,
|
||||
spec_info=None,
|
||||
seq_lens_cpu=new_seq_lens.cpu(),
|
||||
spec_info=None,
|
||||
)
|
||||
backend.init_forward_metadata_out_graph(replay_fb)
|
||||
|
||||
# Verify replay updated the metadata
|
||||
replay_metadata = backend.forward_decode_metadata
|
||||
|
||||
Reference in New Issue
Block a user